Original: Learning Robust Dynamics through Variational Sparse Gating
Revised: Learning Robust Dynamics through Variational Sparse Gating

Original: Anonymous Author(s)

Original: Affiliation

Original: Address email

Original: Abstract
Revised: Abstract

Original: Latent dynamics models learn an abstract representation of the environment basedon collected experience. For example, world models can imagine unseen trajectories, potentially improving sample efficiency in model-based reinforcementlearning.
Revised: 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.

Original: Planning in the real-world requires agents to understand long-termdependencies between actions and events, and account for a varying degree ofchanges, e.g. due to a change in background or viewpoint. These changes areoften quite sparse which suggests incorporating such an inductive bias in a dynamics model.
Revised: In environments with many objects, often only a small number of them are moving or interacting at the same time.

Original: In this work, we introduce Variational Sparse Gating (VSG), wheremodel states are sparselyupdated througha stochastic gating mechanism. Moreover, latent state in prior world models comprise of a deterministic and stochasticpath and they complement each other for solving tasks.
Revised: 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.

Original: We also propose  Simple Variational Sparse Gating (SVSG), which has a purely stochastic latent state andleverages the gating mechanism proposed in VSG.
Revised: 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.

Original: Finally, to evaluate agentsin partial-observability and stochasticity, we also introduce a novel environment,called BringBackShapes (BBS). We conducted experiments on BBS and existingbenchmarks to demonstrate the benefits of proposed methods.
Revised: 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.

Original:  Introduction
Revised:  Introduction

Original: Latent dynamics models are models that generate agent’s future states in the compact latent spacewithout feeding the high-dimensional observations back to the model.
Revised: 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.

Original: 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)
Revised: 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)

Original: (Hafner et al., 2020; 2021; 2019; Ha and Schmidhuber,2018) and robotics (Watter et al., 2015).
Revised: (Hafner et al., 2020; 2021; 2019; Ha and Schmidhuber, 2018) and robotics (Watter et al., 2015).

Original: Generating high dimensional inputs like images in latentspace reduces the accumulating errors leading to more accurate long-term predictions.
Revised: Generating sequences in the compact latent space reduces the accumulating errors leading to more accurate long-term predictions.

Original: Additionally,having lower dimensionality leads to a lower memory footprint.
Revised: Additionally, having lower dimensionality leads to a lower memory footprint.

Original: Solving tasks in model-based RLinvolves learning a world model (Ha and Schmidhuber, 2018) that can predict outcomes of actions,followed by using them to derive behaviors (Sutton, 1991).
Revised: Solving tasks in model-based RL involves learning a world model (Ha and Schmidhuber, 2018) that can predict outcomes of actions, followed by using them to derive behaviors (Sutton, 1991).

Original: Motivated by these benefits, the recentlyproposed DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021) agents achievedstate-of-the-art results on a wide range of visual control tasks.
Revised: Motivated by these benefits, the recently proposed DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021) agents achieved state-of-the-art results on a wide range of visual control tasks.

Original: Many complex tasks require reliable long-term prediction of dynamics.
Revised: Many complex tasks require reliable long-term prediction of dynamics.

Original: This is true especially inpartially observable environments where only a subspace is visible to the agent, and it is usuallyrequired to accurately retain information over multiple time steps to solve the task.
Revised: This is true especially in partially observable environments where only a subspace is visible to the agent, and it is usually required to accurately retain information over multiple time steps to solve the task.

Original: The Dreameragents (Hafner et al., 2020; 2021) employ an Recurrent State-Space Model (RSSM)
Revised: The Dreamer agents (Hafner et al., 2020; 2021) employ an Recurrent State-Space Model (RSSM)

Original: (Hafner et al. comprising of a Recurrent Neural Network (RNN).
Revised: (Hafner et al., 2019) comprising of a Recurrent Neural Network (RNN).

Original: Training RNNs for long sequences is challenging as they suffer from optimization problems like vanishing gradients (Hochreiter, 1991; Bengioet al., 1994).
Revised: Training RNNs for long sequences is challenging as they suffer from optimization problems like vanishing gradients (Hochreiter, 1991; Bengio et al., 1994).

Original: Different ways of applying sparse updates in RNNs have been investigated (Camposet al., 2017; Neil et al., 2016; Goyal et al., 2019), enabling a subset of state dimensions to be constantduring the update.
Revised: Different ways of applying sparse updates in RNNs have been investigated (Campos 36th Conference on Neural Information Processing Systems (NeurIPS 2022). et al., 2017; Neil et al., 2016; Goyal et al., 2019), enabling a subset of state dimensions to be constant during the update.

Original: A sparse update prior is also  motivated by the fact that in the real world, many factors of variation are constant over extended periods of time.
Revised: 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.

Original: For instance, several objects in aphysical simulation may be stationary until some force acts upon them.
Revised: For instance, several objects in a physical simulation may be stationary until some force acts upon them.

Original: Additionally, this is useful inthe partially observable setting where the agent observes a constrained viewpoint and has to keeptrack of objects that are not visible for many time steps.
Revised: Additionally, this is useful in the partially observable setting where the agent observes a constrained viewpoint and has to keep track of objects that are not visible for many time steps.

Original: In this work, we introduce Variational Sparse
Revised: In this work, we introduce Variational Sparse

Original: Gating (VSG), a stochastic gating mechanism that sparsely updates the latent states at each step.
Revised: Gating (VSG), a stochastic gating mechanism that sparsely updates the latent states at each step.

Original: Recurrent State-Space Model (RSSM) (Hafner et al., 2019) was introduced in PLaNet where themodel state was composed of two paths, a image representation path and a recurrent path.
Revised: 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 a recurrent path.

Original: DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021) utilized them to achieve stateof-the-art results in continuous and discrete control tasks (Hafner et al., 2019).
Revised: DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021) utilized them to achieve stateof-the-art results in continuous and discrete control tasks (Hafner et al., 2019).

Original: While the imagerepresentation path which is stochastic accounts for multiple possible future states, the recurrentpath which is deterministic is motivated by stable retention of information over multiple time stepsto facilitate gradient-based optimization.
Revised: 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 optimization.

Original: The authors in PLaNet (Hafner et al., 2019) showed thatboth components were important for solving tasks, where the stochastic part was more important toaccount for partial observability of the initial states.
Revised: (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.

Original: By leveraging the proposed gating mechanism(Variational Sparse Gating (VSG)), we demonstrate that a purely stochastic model with a singlecomponent can achieve competitive results, and call it Simple Variational Sparse Gating (SVSG).
Revised: 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).

Original: To the best of our knowledge, this is the first work that shows that purely stochastic models achievecompetitive performance on continuous control tasks when compared to leading agents.
Revised: 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.

Original: Existing benchmarks (Bellemare et al., 2013; Chevalier-Boisvert et al., 2018; Tassa et al., 2018) for RLdo not test the capability of agents in both partial observability and stochasticity.
Revised: 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.

Original: The Atari (Bellemareet al., 2013) benchmark comprises of 55 games but most of the games are deterministic and a lot ofcompute is required to train on them.
Revised: 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.

Original: Some tasks in the Atari and Minigrid benchmarks are partiallyobservable but either lack stochasticity or are hard exploration tasks.
Revised: Some tasks in the Atari and Minigrid benchmarks are partiallyobservable but either lack stochasticity or are hard exploration tasks.

Original: Also, these benchmarks do notallow for controlling the factors of variation.
Revised: Also, these benchmarks do not allow for controlling the factors of variation.

Original: We developed a new partially-observable and stochasticenvironment, called BringBackShapes (BBS), where the task is to push objects to a predefined goalarea.
Revised: We developed a new partially-observable and stochastic environment, called BringBackShapes (BBS), where the task is to push objects to a predefined goal area.

Original: Solving tasks in BBS require agents to remember states of previously observed objects andavoid noisy distractor objects.
Revised: Solving tasks in BBS require agents to remember states of previously observed objects and avoid noisy distractor objects.

Original: Furthermore, VSG and SVSG outperformed leading model-based andmodel-free baselines.
Revised: Furthermore, VSG and SVSG outperformed leading model-based and model-free baselines.

Original: We also present studies with varying partial-observability and stochasticity todemonstrate that the proposed agents have better memory for tracking observed objects and are morerobust to increasing levels of noise.
Revised: We also present studies with varying partial-observability and stochasticity to demonstrate that the proposed agents have better memory for tracking observed objects and are more robust to increasing levels of noise.

Original: Lastly, the proposed methods were also evaluated on existingbenchmarks - DeepMind Control (DMC) (Tassa et al., 2018), DMC with Natural Background (Zhanget al., 2021; Nguyen et al., 2021b), and Atari (Bellemare et al., 2013).
Revised: Lastly, the proposed methods were also evaluated on existing benchmarks - DeepMind Control (DMC) (Tassa et al., 2018), DMC with Natural Background (Zhang et al., 2021; Nguyen et al., 2021b), and Atari (Bellemare et al., 2013).

Original: On the existing benchmarks,the proposed method performed better on tasks with changing viewpoints and sparse rewards.
Revised: On the existing benchmarks, the proposed method performed better on tasks with changing viewpoints and sparse rewards.

Original: Our key contributions are summarized as follows:
Revised: Our key contributions are summarized as follows:

Original: • Variational Sparse Gating : We introduce Variational Sparse Gating (VSG), where the recurrentstates are sparsely updated through a stochastic gating mechanism.
Revised: • Variational Sparse Gating : We introduce Variational Sparse Gating (VSG), where the recurrent states are sparsely updated through a stochastic gating mechanism.

Original: A comprehensive empiricalevaluation shows that VSG outperforms baselines on tasks requiring long-term memory.•
Revised: A comprehensive empirical evaluation shows that VSG outperforms baselines on tasks requiring long-term memory.

Original: Simple Variational Sparse Gating : We also propose Simple Variational Sparse Gating (SVSG)which has a purely stochastic state, and achieves competitive results on continuous control taskswhen compared with agents that also use a deterministic component.• BringBackShapes : We developed the BringBackShapes (BBS) environment to evaluate agentson partially-observable and stochastic settings where these variations can be controlled.
Revised: Simple Variational Sparse Gating : We also propose Simple Variational Sparse Gating (SVSG) which has a purely stochastic state, and achieves competitive results on continuous control tasks when compared with agents that also use a deterministic component. • BringBackShapes : We developed the BringBackShapes (BBS) environment to evaluate agents on partially-observable and stochastic settings where these variations can be controlled.

Original: Ourexperiments show that the proposed agents are more robustwhen exposed to such variations.
Revised: Our experiments show that the proposed agents are more robust to such variations.

Original:  Variational Sparse Gating
Revised:  Variational Sparse Gating

Original: Reinforcement Learning : The visual control task can be formulated as a Partially Observable
Revised: Reinforcement Learning : The visual control task can be formulated as a Partially Observable

Original: Markov Decision Process (POMDP) with discrete time steps t ∈ [1; T ] .
Revised: Markov Decision Process (POMDP) with discrete time steps t ∈ [1; T ] .

Original: The agent selects actiona t ∼ p ( a t | o ≤ t , a <t ) to interact with the environment and receives the next observation and scalarreward o t , r t ∼ p ( o t , r t | o <t , r <t ), respectively, at each time step.
Revised: The agent selects action a t ∼ p ( a t | o ≤ t , a <t ) to interact with the environment and receives the next observation and scalar reward o t , r t ∼ p ( o t , r t | o <t , r <t ), respectively, at each time step.

Original: The goal is to learn a policy thatmaximizes the expected discounted sum of rewards E p ( (cid:80) T t =1 γ t r t ) , where γ is the discount factor.
Revised: The goal is to learn a policy that maximizes the expected discounted sum of rewards E p ( (cid:80) Tt =1 γ t r t ) , where γ is the discount factor.

Original: Agent : Agent is composed of a world model and a policy (Fig. 1).
Revised: Agent : Agent is composed of a world model and a policy (Fig. 1).

Original: World models (Sec. 2.1) encodea sequence of observations and actions into latent representations.
Revised: World models (Sec. 2.1) encode a sequence of observations and actions into latent representations.

Original: Behaviors (Appendix B) arederived to maximize expected returns on the trajectories generated from the learned world model.
Revised: The agents behavior (Appendix B) is derived to maximize expected returns on the trajectories generated from the learned world model.

Original: While training, the world model is learned with collected experience, the policy is improved ontrajectories unrolled using the world model and new episodes are collected by deploying the policy 𝑠𝑠𝑠 3 𝑠𝑎Ƹ𝑟𝑥𝑥𝑥ො𝑥𝑎𝑎Ƹ𝑟ො𝑥
Revised: 𝑠𝑠𝑠 3 𝑠𝑎Ƹ𝑟𝑥𝑥𝑥ො𝑥𝑎𝑎Ƹ𝑟ො𝑥

Original: VSG VSG VSG Ƹ𝑟ො𝑥(a)
Revised: VSG VSG VSG Ƹ𝑟ො𝑥(a)

Original: 𝑠 1 𝑠Ƹ𝑟𝑥Ƹ𝑟
Revised: 𝑠 1 𝑠Ƹ𝑟𝑥Ƹ𝑟

Original: VSG VSG VSG Ƹ𝑟 3 ො𝑎ො𝑎ො𝑎ො𝑣ො𝑣ො𝑣Ƹ𝑠Ƹ𝑠(b)in the environment.
Revised: VSG VSG VSG Ƹ𝑟 3 ො𝑎ො𝑎ො𝑎ො𝑣ො𝑣ො𝑣Ƹ𝑠Ƹ𝑠(b)

Revised: While training, the world model is learned with collected experience, the policy is improved on trajectories unrolled using the world model and new episodes are collected by deploying the policy in the environment.

Original: An initial set of episodes are collected using a random policy.
Revised: An initial set of episodes are collected using a random policy.

Original: As trainingprogresses, new episodes are collected using the latest policy to further improve the world model.
Revised: As training progresses, new episodes are collected using the latest policy to further improve the world model.

Original: 2.1 World Model
Revised: 2.1 World Model

Original: World Models (Ha and Schmidhuber, 2018) learn to mimic the environment, using the collectedexperience, to facilitate deriving behaviours in the abstract latent space.
Revised: World Models (Ha and Schmidhuber, 2018) learn to mimic the environment using the collected experience and facilitate deriving behaviours in the abstract latent space.

Original: Given an abstract state of theworld and an action, the model applies the learned transition dynamics to predict the resulting nextstate and reward.
Revised: Given an abstract state of the world and an action, the model applies the learned transition dynamics to predict the resulting next state and reward.

Original: RSSM (Hafner et al., 2019) was introduced in PlaNet, where the model state wascomposed of two paths.
Revised: RSSM (Hafner et al., 2019) was introduced in PlaNet, where the model state was composed of two paths.

Original: The recurrent path consists of an RNN (See Figure 2 [a]), and is motivatedwith reliable long-term information preservation, while the image representation path samples from alearned distribution to account for multiple possible futures (Babaeizadeh et al., 2017).
Revised: The recurrent path consists of an RNN (See Figure 2 [a]), and is motivated with reliable long-term information preservation, while the image representation path samples from a learned distribution to account for multiple possible futures (Babaeizadeh et al., 2017).

Original: In this work,we introduce Variational Sparse Gating (VSG), where the recurrent path selectively updates a subsetof the latent states at each step using a stochastic gating network.
Revised: In this work, we introduce Variational Sparse Gating (VSG), where the recurrent path selectively updates a subset of the latent states at each step using a stochastic gating network.

Original: Sparse updates enables agents tohave long-term memory and learn robust representations which is required for complex tasks.
Revised: Sparse updates enable the agent to have long-term memory and learn robust representations to solve complex tasks.

Original: Model Components : The world model comprises of an image encoder, a VSG model, and predictorsfor image, discount and reward.
Revised: Model Components : The world model comprises of an image encoder, a VSG model, and predictors for image, discount and reward.

Original: The image encoder generates representations o t for the observation x tusing CNNs.
Revised: The image encoder generates representations o t for the observation x t using Convolutional Neural Networks (CNNs).

Original: The VSG model comprises of a recurrent network equipped with the stochastic gatingmechanism, and is used to compute two stochastic image representation states.
Revised: The VSG model comprises of a recurrent model equipped with the stochastic gating mechanism to get the recurrent state h t , and is used to compute two stochastic image representation states.

Original: The posteriorrepresentation state z t  contains information about the current image x t , while the prior state ˆ z tdoes not observe x t .
Revised: The posterior representation state z t is obtained using the representation model and contains information about the current observation x t . The prior state ˆ z t is obtained from the transition predictor without observing the current observation x t .

Original: For planning, sequences are generated in compact latent state by utilizing theprior representation state ˆ z t .
Revised: This is useful while planning as sequences are generated in compact latent state, and the output from the transition predictor is utilized.

Original: This  results in a lower memory footprint and enables predictions ofthousands of trajectories in parallel on a single GPU.
Revised: This also results in a lower memory footprint and enables predictions of thousands of trajectories in parallel on a single GPU.

Original: The representation states are sampled from aknown distribution with learned parameters like Gaussian (Hafner et al., 2020) or Categorical (Hafneret al., 2021).
Revised: The representation states are sampled from a known distribution with learned parameters like Gaussian (Hafner et al., 2020) or Categorical (Hafner et al., 2021).

Original: The concatenation of outputs from the recurrent and image representation models givesthe compact model state ( s t =
Revised: The concatenation of outputs from the recurrent and image representation models gives the compact model state ( s t =

Original: [ h t , z t ] ).
Revised: [ h t , z t ] ).

Original: The posterior model state is further used to reconstruct theoriginal image ˆ x t , predict the reward ˆ r t , and discount factor ˆ γ t .
Revised: The posterior model state is further used to reconstruct the original image ˆ x t , predict the reward ˆ r t , and discount factor ˆ γ t .

Original: The discount factor helps to predictthe probability that an episode will end.
Revised: The discount factor helps to predict the probability that an episode will end.

Original: The components of  world model are as follows:
Revised: The components of the world model are as follows:

Original: Recurrent model:

Original: [Equation]
Revised: [Equation]

Original: Representation [Equation] [Equation] predictor
Revised: 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝑢 𝑡 ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊

Original: [Equation] Reward [Equation] Discount predictor
Revised: 𝑢 𝜎 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 𝑠 𝑡−

Original: [Equation] Neural Networks : The representation model outputs the posterior image representation state z tconditioned on the image encoding x t and recurrent state h t . The transition predictor provides theprior image representation state ˆ z t .
Revised: 𝑊 𝑣 ⊗ 𝑖 𝑡 𝑣 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 ℳ ℳ ǁ𝑠 𝑡 𝑠 𝑡 Ƹ𝑠 𝑡 (a) 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝑢

Original: The image encoding o t is obtained by passing the image x t

Original: 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝑢 𝑡 ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 𝑠 𝑡−
Revised: 𝑡 ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 𝑠 𝑡−

Original: 𝑊 𝑣 ⊗ 𝑖 𝑡 𝑣 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 ℳ ℳ ǁ𝑠 𝑡 𝑠 𝑡 Ƹ𝑠 𝑡 (a) 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝑢
Revised: 𝑊 𝑣 ⊗ 𝑖 𝑡 𝑣 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 ℳ ℳ ǁ𝑠 𝑡 𝑠 𝑡 Ƹ𝑠 𝑡 (b) 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝑢

Original: 𝑡 ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 𝑠 𝑡−
Revised: 𝑡 ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 𝑠 𝑡−

Original: 𝑊 𝑣 ⊗ 𝑖 𝑡 𝑣 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 ℳ ℳ ǁ𝑠 𝑡 𝑠 𝑡 Ƹ𝑠 𝑡 (b) 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝑢
Revised: 𝑊 𝑣 ⊗ 𝑖 𝑡 𝑣 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 ℳ ℳ ǁ𝑠 𝑡 𝑠 𝑡 Ƹ𝑠 𝑡 (c)

Original: 𝑡 ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 𝑊 𝑣 ⊗ 𝑖 𝑡 ℎ 𝑡− ℎ 𝑡 𝑣 𝑡 ෨ℎ 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ ℳ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 𝑠 𝑡−
Revised: Neural Networks : The representation model outputs the posterior image representation state z t conditioned on the image encoding x t and recurrent state h t .

Original: 𝑊 𝑣 ⊗ 𝑖 𝑡 𝑣 𝑡 𝑜 𝑡 𝑧 𝑡 Ƹ𝑧 𝑡 ~ ~ tanh ⊕ 𝜎 𝑓 𝑝 𝑓 𝑞 𝑊 𝑐 𝑊 𝑢 𝜎 ~ℬ ෤𝑢 𝑡 𝑢 𝑡 ℳ ℳ ǁ𝑠 𝑡 𝑠 𝑡 Ƹ𝑠 𝑡 (c) through CNN (LeCun et al., 1989) and MLP layers.
Revised: The transition predictor provides the prior image representation state ˆ z t . The image encoding o t is obtained by passing the image x t through CNN (LeCun et al., 1989) and Multi-layer Perceptron (MLP) layers.

Original: In VSG, we propose to modify the Gated Recurrent Unit (GRU) used in RSSM to sparsely update the recurrent state at each step.
Revised: In VSG, we propose to modify the Gated Recurrent Unit (GRU) used in RSSM to sparsely update the recurrent state at each step.

Original: The modelstate s t , which is a concatenation of recurrent and image representation states is passed throughseveral layers of MLP to predict the discount and reward, and transposed CNN layers are used toreconstruct the image.
Revised: The model state s t , which is a concatenation of recurrent and image representation states is passed through several layers of MLP to predict the discount and reward, and transposed CNN layers are used to reconstruct the image.

Original: The Exponential Linear Unit (ELU) activation is used for training all thecomponents of the world model (Clevert et al., 2015).
Revised: The Exponential Linear Unit (ELU) activation is used for training all the components of the world model (Clevert et al., 2015).

Original: Sparse Gating : In light of training RNNs to capture long-term dependencies, different ways ofapplying sparse updates have been investigated (Campos et al., 2017; Neil et al., 2016; Goyalet al., 2019), enabling a subset of state dimensions to be constant during the update.
Revised: Sparse Gating : In light of training RNNs to capture long-term dependencies, different ways of applying sparse updates have been investigated (Campos et al., 2017; Neil et al., 2016; Goyal et al., 2019), enabling a subset of state dimensions to be constant during the update.

Original: They werefound to alleviate the vanishing gradient problem by effectively reducing the number of sequentialoperations (Campos et al., 2017).
Revised: They were found to alleviate the vanishing gradient problem by effectively reducing the number of sequential operations (Campos et al., 2017).

Original: Discrete gates may also improve long-term memory by avoidingthe gradual change of state values introduced by repeated multiplication with continuous gate valuesin standard recurrent architectures.
Revised: 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.

Original: Previous works on sparsely updating hidden states (Campos et al. ,2017; Neil et al., 2016) use a separate layer applied over the outputs of RNN, and do not modify the
Revised: Previous works on sparsely updating hidden states (Campos et al., 2017; Neil et al., 2016) use a separate layer applied over the outputs of RNN, and do not modify the

Original: RNN in itself.
Revised: RNN in itself.

Original: However, in this work, we modify the update gate in GRU (Cho et al., 2014) to takebinary values by sampling from a Bernoulli distribution (Fig. 

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

Original: [b] shows the architecture).
Revised: [b] shows the architecture).

Original: The input i t to the recurrent model contains information about the action and is obtained by concatenating the previous image representation state z t − 1 and action a t followed by passing them throughthrough a MLP layer.
Revised: The input i t to the recurrent model contains information about the action and is obtained by concatenating the previous image representation state z t − 1 and action a t followed by passing them through a MLP layer.

Original: Similar to GRU (Cho et al., 2014), there are two gates- reset and update.
Revised: Similar to GRU (Cho et al., 2014), there is a reset and update gate.

Original: Thereset gate v t decides the extent of information flow from the previous recurrent state and inputs, andthe update gate u t tells which parts of the recurrent state will be updated.
Revised: The reset gate v t decides the extent of information flow from the previous recurrent state and inputs, and the update gate u t tells which parts of the recurrent state will be updated.

Original: The update gate takes onlybinary values, selecting whether the value will be updated or copied from previous time step.
Revised: The update gate takes only binary values, selecting whether the value will be updated or copied from previous time step.

Original: Binaryvalues are obtained by sampling from a Bernoulli distribution where the probability of sampling isobtained using the previous recurrent state h t − 1 and input i t .
Revised: Binary values are obtained by sampling from a Bernoulli distribution where the probability of sampling is obtained using the previous recurrent state h t − 1 and input i t .

Original: Straight-through estimators (Bengioet al., 2013) were used for propagating gradients backwards for training.
Revised: Straight-through estimators (Bengio et al., 2013) were used for propagating gradients backwards for training.

Original: The update equations are asfollows:
Revised: The update equations are:

Original: [Equation]
Revised: [Equation]

Original: where ⊙ denotes element-wise multiplication, σ and tanh are the sigmoid and hyperbolic tangentactivation function, and W ∗ and b ∗ denotes the weights and biases, respectively.
Revised: where ⊙ denotes element-wise multiplication, σ and tanh are the sigmoid and hyperbolic tangent activation function, and W ∗ and b ∗ denotes the weights and biases, respectively.

Original: To control thesparsity of updates, we have used KL divergence between probability of sampling the update gate ˜ u tand a fixed prior probability κ , where κ is a tunable hyperparameter.
Revised: To control the sparsity of updates, we have used KL divergence between probability of sampling the update gate ˜ u t and a fixed prior probability κ , where κ is a tunable hyperparameter.

Original: Loss function : The predictors for image and reward produces Gaussian distributions with unitvariance, whereas the discount predictor predicts a Bernoulli likelihood.
Revised: Loss function : The predictors for image and reward produces Gaussian distributions with unit variance, whereas the discount predictor predicts a Bernoulli likelihood.

Original: The image representationstates are sampled from a Gaussian (Hafner et al., 2020) or a Categorical (Hafner et al., 2021)distribution which are trained to maximize the likelihood of targets.
Revised: The image representation states are sampled from a Gaussian (Hafner et al., 2020) or a Categorical (Hafner et al., 2021) distribution which are trained to maximize the likelihood of targets.

Original: In addition, there is a KL
Revised: In addition, there is a KL

Original: Divergence term between the prior and posterior distributions and similar to DreamerV2 (Hafneret al., 2021), we have also used KL balancing with a factor of 0.8.
Revised: Divergence term between the prior and posterior distributions and similar to DreamerV2 (Hafner et al., 2021), we have also used KL balancing with a factor of 0.8.

Original: We have also added a sparsity lossto regularize the number of updates in hidden state at each step.
Revised: We have also added a sparsity loss to regularize the number of updates in hidden state at each step.

Original: All the components of the worldmodel are optimized jointly using the loss function given by:
Revised: All the components of the world model are optimized jointly using the loss function given by:

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: where β and α are the scale for KL losses of the latent codes and the sparse update gates, respectively.
Revised: where β and α are the scale for KL losses of the latent codes and the sparse update gates, respectively.

Original:  Simple Variational Sparse Gating
Revised:  Simple Variational Sparse Gating

Original: Stochastic State-Space Model (SSM) were proposed in PLaNet (Hafner et al., 2019), where it wasdiscussed that it is not trivial to achieve competitive results without the deterministic recurrent path.
Revised: 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.

Original: Having a deterministic component was motivated to allow the transition model to retain informationfor multiple time steps as the stochastic component induces variance (Hafner et al., 2019).
Revised: Having a 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).

Original: In thiswork, we show that having a purely stochastic component achieves comparable performance with
Revised: In this work, we show that having a purely stochastic component achieves comparable performance with

Original: DreamerV2 while significantly outperforming SSMs (refer to Appendix H for more details).
Revised: DreamerV2 while significantly outperforming SSMs (refer to Appendix H for more details).

Original: Weintroduce a simplified version of VSG, called Simple Variational Sparse Gating (SVSG) where theworld model has a model states t with single path to preserve information over multiple steps andalso account for partial observability in future states (Fig. {+
+}
Revised: We introduce a 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. {+
+}

Original: [c] presents the SVSG architecture).
Revised: [c] presents the SVSG architecture).

Original: Themodel state s t is sparsely updated at each step using the stochastic gating mechanism proposed in VSG. In SVSG,  the posterior state s t is conditioned on the previous state s t − 1 , input image x t andthe action a t , whereas the prior state ˆ s t is generated without using the  image observation.
Revised: Model Components : In SVSG, there is no recurrent model and the posterior state s t is obtained using the representation model by conditioning on the previous state s t − 1 , input image x t and the action a t . Similar to VSG, there is a transition predictor that returns the prior state ˆ s t which does not use the current image observation to imagine trajectories in the latent space.

Revised: Both the modules sparsely update the model state at each step using the stochastic gating mechanism proposed in VSG.

Original: We haveused adiagonal Gaussian distribution for the stochastic state with a learnable mean  and standarddeviation.
Revised: We have used a Gaussian distribution for the stochastic state with a learnable mean vector and a learnable diagonal covariance matrix.

Original: Similar to VSG, the posterior state is used to reconstruct the image, and predict the rewardand discount factor.
Revised: Similar to VSG, the posterior state is used to reconstruct the image, and predict the reward and discount factor.

Original: The components of world models in SVSG are:
Revised: The components of world model in SVSG are:

Original: Representation model
Revised: Representation model:

Revised: [Equation]

Revised: [Equation]

Revised: Reward predictor

Original: [Equation]
Revised: [Equation]

Original: The representation model q ϕ and transition predictor p ϕ are modified to output the posterior s t andprior ˆ s t states, respectively.
Revised: The representation model q ϕ and transition predictor p ϕ are modified to output the posterior s t and prior ˆ s t states, respectively.

Revised: The reset gate v t and the update gate ˜ u t is calculated using the previous state s t − 1 and input i t which has the information about the action a t .

Original: The candidate state ˜ s t at each step is obtained using input i t , reset gate v tand previous state s t − 1 .
Revised: The candidate state ˜ s t at each step is obtained using input i t , reset gate v t and previous state s t − 1 .

Original: Similar to VSG, the update gate u t is sampled from a Bernoulli distributionto sparsely update the latent states at each step, given by:[Equation]
Revised: Similar to VSG, the update gate u t is sampled from a Bernoulli distribution to sparsely update the latent states at each step, given by:

Original: [Equation]
Revised: [Equation]

Original: where ⊙ denotes the element-wise multiplication, σ and tanh are the sigmoid and hyperbolic tangentactivation functions, and W ∗ and b ∗ denote the weights and biases, respectively.
Revised: where ⊙ denotes the element-wise multiplication, σ and tanh are the sigmoid and hyperbolic tangent activation functions, and W ∗ and b ∗ denote the weights and biases, respectively.

Original: The candidate state ˜ s t was passed through MLP layers to get the prior and posterior distributions.
Revised: The candidate state ˜ s t is feeded through MLP layers to get the prior and posterior distributions.

Original: The image encoding x t was used to get posterior distribution, whereas the prior distribution waspredicted without it.
Revised: The image encoding x t was used to get posterior distribution, whereas the prior distribution was predicted without it.

Original: After sampling the prior ˆ z t and posterior z t candidate states, the update gatesparsely modifies the previous latent state and outputs the prior ˆ s t and posterior s t model state ateach step, respectively.
Revised: The prior ˆ z t and posterior z t candidate states are sampled from these distributions, where the update gate sparsely modifies the previous latent state and outputs the prior ˆ s t and posterior s t model states at each step, respectively.

Original: The update equations are given by:
Revised: The update equations are given by:

Original: [Equation]
Revised: [Equation]

Original: where f p and f q denotes functions that output a distribution with learnable parameters for prior andposterior, respectively.
Revised: where f p and f q denotes functions that output a distribution with learnable parameters for prior and posterior, respectively.

Original: ∼ denotes sampling from the learned distribution.

Original: For SVSG, Categoricallatents (Hafner et al., 2021) were not performing well on our tasks.
Revised: For SVSG, Categorical latents (Hafner et al., 2021) were not performing well on our tasks.

Original: We attribute this to the fact thatsamples from a categorical distribution are binary vectors and it is difficult to accurately reconstructwith such sparse latent representations.
Revised: 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.

Original: Lastly, we observed thatthe sparse gating mechanismintroduced in VSG was important forthe convergence of SVSG.
Revised: Lastly, we observed that sparse gating mechanism introduced in VSG was important for convergence of SVSG.

Revised: Loss function : We have used the same loss function as described in Sec. 2 and policy is similar to used in VSG (described in Appendix B).

Original: For training the SVSG model, we replace the KL loss term between prior and posterior distributions in Eq.
Revised: For training the SVSG model, we replace the KL loss term between prior and posterior distributions in Eq.

Original: 3 with a masked KL loss that penalizes the state dimensions that were updated in the correspondingtime step, i.e. those for which the corresponding element in u t is equal to 1.
Revised: 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 u t is equal to 1.

Original: We found this to benecessary, since the original, unmasked KL loss did not yield good performance, presumably due toits effect on state dimensions that were not updated.
Revised: We found this to be necessary, since the original, unmasked KL loss did not yield good performance, presumably due to its effect on state dimensions that were not updated.

Original: We have used the same loss function as describedin Sec. 2 and policy is similar to used in VSG (described in Appendix B).

Original:  Experiments 
Revised:  Experiments 4.1 BringBackShapes

Original: 4.1 BringBackShapes

Original:  In this work, we developed the BringBackShapes (BBS) environment to test the ability of agentsto solve tasks in partially-observable and stochastic scenarios (see Fig.
Revised: Environment : In this work, we developed the BringBackShapes (BBS) environment to test the ability of agents to solve tasks in partially-observable and stochastic scenarios (see Fig.

Original: 3 [a]).
Revised: 3 [a]).

Original: The task is to pushthe objects within the arena into a pre-specified goal area.
Revised: The task is to push the objects within the arena into a pre-specified goal area.

Original: Moreover, rewards are sparse and is+1 for successfully pushing an object into the goal.
Revised: Moreover, rewards are sparse and is +1 for successfully pushing an object into the goal.

Original: At each time step, the agent only receives anobfuscated view of the arena centered around its current position.
Revised: At each time step, the agent only receives an obfuscated view of the arena centered around its current position.

Original: This requires agents to efficientlyexplore the arena to find new objects as well as remember states of previously observed objects.
Revised: This requires agents to efficiently explore the arena to find new objects as well as remember states of previously observed objects.

Original: The objects can collide with each other and the walls, which further requires the agent to accountfor these events while updating its state.
Revised: The objects can collide with each other and the walls, which further requires the agent to account for these events while updating its state.

Original: Stochasticity was introduced in the environment by usingrandom distractor objects.
Revised: Stochasticity was introduced in the environment by [Table] [Figure] [Table]

Revised: VSG and SVSG significantly outperformed DreamerV2 on using random distractor objects.

Original: The distractor objects follow Brownian motion in any direction and can’tbe pushed into the goal area.
Revised: The distractor objects follow Brownian motion in any direction and can’t be pushed into the goal area.

Original: Additionally, they add noise to the reward signal as they might pushobjects into the goal, causing a reward that is not or only partially related to the agent’s behavior.
Revised: Additionally, they add noise to the reward signal as they might push objects into the goal, causing a reward that is not or only partially related to the agent’s behavior.

Original: Dueto partial-observability, such instances might not be visible to the agent, making the task even morechallenging.
Revised: Due to partial-observability, such instances might not be visible to the agent, making the task even more challenging.

Original: Refer to Appendix A for further description of the environment.
Revised: Refer to Appendix A for further description of the environment.

Original: Quantitative Results : Fig.
Revised: Experimental Setup : The BringBackShapes (BBS) environment returns high dimensional images of shape 64 × 64 × 3 as observation. Action is a 2-dimensional continuous vector with acceleration and direction as components. Episodes last for 3000 environment steps and an action repeat (Mnih et al., 2016) of 4 was used. Baseline agents include DreamerV2 (Hafner et al., 2021) and DrQ-v2 (Yarats et al., 2022).

Revised: In Appendix C, we mention the hyperparameters for the proposed methods- VSG and

Revised: SVSG.

Revised: The model was implemented using Tensorflow Probabability (Dillon et al., 2017) and trained on a single NVIDIA V100 GPU with 16GB memory.

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

Revised: Lastly, results are reported across 5 seeds.

Revised: Quantitative Results : Fig.

Original: 3 [b] compares the proposed methods VSG and SVSG with the leading
Revised: 3 [b] compares the proposed methods VSG and SVSG with the leading

Original: RL agents- DreamerV2 (Hafner et al., 2021) and DrQ-V2 (Yarats et al., 2022).
Revised: RL agents- DreamerV2 (Hafner et al., 2021) and DrQ-V2 (Yarats et al., 2022).

Original: The score indicateshow many objects on average were scored within a episode.
Revised: The score indicates how many objects on average were scored within a episode.

Original: As discussed in Section ?? , we 
Revised: As discussed in Section 3, we trained

Original: [Figure]
Revised: SVSG with Gaussian latents only.

Original: [Table] competitive scores when compared to the proposed methods.
Revised: Upon evaluation at 2.5M timesteps, DreamerV2 achieves competitive scores when compared to the proposed methods.

Original: Whereas at 1M steps, DreamerV2does not perform as well as VSG, which has mean score of 4.9 .
Revised: Whereas at 1M steps, DreamerV2 does not perform as well as VSG, which has mean score of 4.9 .

Original: This shows that learning withsparsity priors helps improve the sample efficiency.
Revised: This shows that learning with sparsity priors helps improve the sample efficiency.

Original: Furthermore, performance of SVSG is better than DreamerV2 but similar to VSG, demonstrating that a purely stochastic model can achieve similarperformance.
Revised: Furthermore, performance of SVSG is better than DreamerV but similar to VSG, demonstrating that a purely stochastic model can achieve similar performance.

Original: Varying Partial-Observability and Stochasticy : We also study the effect of partial observabilityand stochasticity.
Revised: Varying Partial-Observability and Stochasticy : We also study the effect of partial observability and stochasticity.

Original: For partial-observability, we increased the size of the arena while reducing theportion visible to the agent.
Revised: For partial-observability, we increased the size of the arena while reducing the portion visible to the agent.

Original: We consider 4 configurations of the arena- Basic, Small, Medium and
Revised: We consider 4 configurations of the arena- Basic, Small, Medium and

Original: Large.
Revised: Large.

Original: For stochasticity, we increase the number of distractor objects using values 0, 1, 2, and 4.
Revised: For stochasticity, we increase the number of distractor objects using values 0, 1, 2, and 4.

Original: Fig.
Revised: Fig.

Original: 3 [c] presents the plots of models trained at 2.5M steps at different sizes of the arena and numberof distractor objects.
Revised: 3 [c] presents the plots of models trained at 2.5M steps at different sizes of the arena and number of distractor objects.

Original: It can be observed that increasing the size of arena makes it harder to scoreobjects.
Revised: It can be observed that increasing the size of arena makes it harder to score objects.

Original: VSG was found to outperform DreamerV2 across all arena sizes.
Revised: VSG was found to outperform DreamerV2 across all arena sizes.

Original: However, SVSG did notperform well on larger arena sizes.
Revised: However, SVSG did not [Figure] perform well on larger arena sizes.

Original: Adding noisy distractor objects led to drop in final performanceof all models.
Revised: Adding noisy distractor objects led to drop in final performance of all models.

Original: But VSG and SVSG still outperformed DreamerV2, indicating that the sparsity priorhelps in ignoring the noisy objects in the arena while solving the task.
Revised: But VSG and SVSG still outperformed DreamerV2, indicating that the sparsity prior helps in ignoring the noisy objects in the arena while solving the task.

Original: Ablation Studies : In this work, we also report statistics to describe the behavior of learned agents.
Revised: Ablation Studies : In this work, we also report statistics to describe the behavior of learned agents.

Original: First-visit time is the number of episode steps taken to visit an object (when object is completelyvisible in the agent’s view).
Revised: 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.

Original: Lower first-visit time for all objects tells that an agent is able to quicklydiscover all the objects in the arena.
Revised: Lower first-visit time indicates that an agent is able to quickly discover all the objects in the arena.

Original: Another metric is Episode Length which denotes the number ofsteps taken by the agent to complete the task.
Revised: Another metric is Episode Length which denotes the number of steps taken by the agent to complete the task.

Original: The maximum of these scores was set to 3000.
Revised: The maximum of these scores was set to 3000.

Original: Wealso report the percentage of objects that were visited within an episode which represents the abilityof agents to explore all parts of the arena to find novel objects.
Revised: We also report the percentage of objects that were visited within an episode which represents the ability of agents to explore all parts of the arena to find novel objects.

Original: Lastly, we report the percentage ofvisited objects that were not scored which tells us that the agent might not be remembering positionsof objects seen previously.
Revised: 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.

Original: Thus, agents might have to explore the arena again to find them leading toan increase in time taken to finish the task.
Revised: Thus, agents might have to explore the arena again to find them leading to an increase in time taken to finish the task.

Original: Table 1 presents the results on different settings of theenvironment and it can be observed that VSG and SVSG outperform DreamerV2 significantly.
Revised: Table 1 presents the results on different settings of the environment and it can be observed that VSG and SVSG outperform DreamerV2 significantly.

Original: Qualitative Results : We also observed the maneuvers taken by the agents to push the objects tothe goal.
Revised: Qualitative Results : We also observed the maneuvers taken by the agents to push the objects to the goal.

Original: The DreamerV2 agent was able to recognize objects and go behind them to push, but[Figure] did not follow a smooth trajectory and was spending more time around an object to push it in thegoal.
Revised: The DreamerV2 agent was able to recognize objects and go behind them to push, but did not follow a smooth trajectory and was spending more time around an object to push it in the goal.

Original: However, the proposed methods had comparatively much smoother trajectories and were moreefficient at pushing objects successfully in the goal area.
Revised: However, our methods showed smoother trajectories and were more efficient at pushing objects successfully in the goal area.

Original: Additionally, the proposed methods learnedto avoid noisy distractor objects whereas the DreamerV2 agent was colliding with them and wastrying to push them to the goal (seeFig.
Revised: Additionally, our methods learned to avoid noisy distractor objects whereas the DreamerV2 agent was colliding with them and was trying to push them to the goal (see

Original:  4 and supplementary material for more videos).
Revised: Fig. 4 and supplementary material for more videos).

Revised: Effect of Sparse Gating :

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

Revised: The sequence of actions is kept fixed across rollouts.

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

Revised: It was observed in Appendix J that learned world model in VSG and

Original:  4.2 DeepMind Control Suite 
Revised: 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 Setup :

Original: The proposed method is evaluated on a few tasks from DeepMind Control Suite (Tassa et al., 2018).
Revised: The proposed method is evaluated on a few tasks from DeepMind Control Suite (Tassa et al., 2018).

Original: Observations for the agents are high dimensional images of shape 64 × 64 × 3 , actions range between1 to 12 dimensions, and episodes last for 1000 steps.
Revised: Observations for the agents are high dimensional images of shape 64 × 64 × 3 , actions range between 1 to 12 dimensions, and episodes last for 1000 steps.

Original: An action repeat (Mnih et al., 2016) of 2 wasused.
Revised: An action repeat (Mnih et al., 2016) of 2 was used.

Original: The model was implemented using Tensorflow Probabability (Dillon et al., 2017) and trainedon a single NVIDIA V100 GPU with 16GB memory in less than 6 hours.
Revised: The model was implemented using Tensorflow Probabability (Dillon et al., 2017) and trained on a single NVIDIA V100 GPU with 16GB memory in less than 6 hours.

Original: The agents were trained for1M environment steps.
Revised: The agents were trained for 1M environment steps.

Original: 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).
Revised: Baselines include DreamerV1 (Hafner et al., , DreamerV2 (Hafner et al., 2021), DrQ-v2 (Yarats et al., 2022), D4PG (Barth-Maron et al., 2018), and A3C (Mnih et al., 2016).

Original: Except A3C, all baselines learn policies from high dimensional pixel inputs.
Revised: Except A3C, all baselines learn policies from high dimensional pixel inputs.

Original: DreamerV2 was trainedusing the implementation provided by the authors.
Revised: DreamerV2 was trained using the implementation provided by the authors.

Original: For other baselines, the metrics provided by therespective authors were used for comparison.
Revised: For other baselines, the metrics provided by the respective authors were used for comparison.

Original: Lastly, returns averaged across 5 seeds were reported.
Revised: Lastly, returns averaged across 5 seeds were reported.

Original:  Figure 5 presentsthe comparison of VSG and SVSG with the aforementioned baselines (Refer to Appendix D for comparison ona few more tasks).
Revised: Results : Figure 5 presents comparison of VSG and SVSG with the baseline agents (Refer to Appendix D for comparison on more tasks).

Original: It can be observed that on most tasks, VSG performscomparable to or better than DreamerV2 (Hafner et al., 2021).
Revised: It can be observed that on most tasks, VSG performs comparable to or better than DreamerV2 (Hafner et al., 2021).

Original: Notably, VSG significantly outperforms DreamerV2 on Quadruped and Finger-Spin tasks.
Revised: Notably, VSG significantly outperforms DreamerV2 on Quadruped and Finger-Spin tasks.

Original: Furthermore,it can be observed that SVSG with apurely stochastic component has similar performance to DreamerV2, outperforming on Finger Spinand Quadruped tasks and performing worse on Cartpole Swingup and Cheetah Run tasks.
Revised: Furthermore, SVSG with a purely stochastic component has similar performance to DreamerV2, outperforming on Finger Spin and Quadruped tasks and performing worse on Cartpole Swingup and Cheetah Run tasks.

Original: In addition,

Original:  SVSG significantly outperforms DreamerV1 on many tasks which also used Gaussian latents and RSSM with multiple paths.
Revised: In addition, SVSG significantly outperforms DreamerV1 on many tasks which also used Gaussian latents and RSSM with multiple paths.

Original: 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).
Revised: 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).

Original:  In BBS, we added noise in the environment by having distractor objects.
Revised: Ablation Studies : In BBS, we added noise in the environment by having distractor objects.

Original: We also experimented withother form of noise where natural videos is used in the background (Zhang et al., 2021;
Revised: We also experimented with other forms of noise where natural videos are used in the background for

Original:  Nguyen et al.,2021b).
Revised: DeepMind Control tasks (Zhang et al., 2021; Nguyen et al., 2021b).

Original: Since reconstruction-free model based RL (Deng et al., 2021; Nguyen et al., 2021a) performbetter than reconstruction based agents (Hafner et al., 2021) in such scenarios, we updated the RSSMblock in DreamerPro (Deng et al., 2021) with VSG and call it VSGPro, and trained it on DMC withnatural background (Refer to Appendix F).
Revised: Since reconstruction-free model based RL (Deng et al., 2021; Nguyen et al., 2021a) perform better than reconstruction based agents (Hafner et al., 2021) in 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).

Original: Lastly, we also compare VSG on discrete control tasksfrom Atari benchmark in Appendix G and VSG performed better on task with changing viewpoints.
Revised: 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.

Original:  Related Work
Revised:  Related Work

Revised: Latent Dynamics Models : Latent dynamics models (Bourlard and Morgan, 2012; Kalman, 1960;

Revised: Bengio et al., 1999) operate directly on sequences predicted in the latent space rather than autoregressively feeding back the generated frames back to the model.

Original: Latent Dynamics Models : Recent advancements in deep learning have allowed learning expressivelatent dynamics models using stochastic backpropagation (Kingma and Welling, 2013; Chung et al. ,2015; Krishnan et al., 2015; Karl et al., 2016).
Revised: 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).

Original: Recurrent State-Space Model (RSSM) (Hafner et al., 2019) comprised of stochastic and deterministic components.
Revised: Recurrent StateSpace Model (RSSM) (Hafner et al., 2019) comprises of stochastic and deterministic components.

Original: VideoFlow (Kumar et al. ,2019) predicted future values of the latent state by normalizing flows for robotic object interactions.
Revised: VideoFlow (Kumar et al., 2019) predicted future values of the latent state by normalizing flows for robotic object interactions.

Original: Hierarchical latent models for video prediction were proposed in CWVAEs (Saxena et al., 2021) withlevels ticking at different intervals in time.
Revised: Hierarchical latent models for video prediction were proposed in CWVAEs (Saxena et al., 2021) with levels ticking at different intervals in time.

Original: (Franceschi et al., 2020; Donà et al., 2021) disentangleddynamic and static factors where 2-5 initial observations was used to estimate the static component.
Revised: (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.

Original: RL for Visual Control : Deep Reinforcement Learning (DRL) methods fall into one of two categories:1) Model-Based — where an explicit model of the environment and its dynamics are learned (Haand Schmidhuber, 2018; Hafner et al., 2019; 2020; 2021), and 2) Model-Free — where a policy islearned directly from the raw observations (Srinivas et al., 2020; 
Revised: RL for Visual Control : Deep Reinforcement 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 et al., 2015; Yarats et al., 2022;

Original: Kostrikov et al., 2020b; Lillicrapet al., 2015; Yarats et al., 2022; Schwarzer et al., 2021).
Revised: 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.

Original: CURL (Srinivas et al., 2020) uses contrastivelosses to learn discriminative representations.
Revised: CURL (Srinivas et al., 2020) uses contrastive losses to learn discriminative representations.

Original: DrQ (Kostrikov et al., 2020a) and DrQ-v2 (Yaratset al., 2022) employed data augmentation techniques and do not use auxiliary losses or pre-training.
Revised: 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.

Original: RSSMs was introduced in PlaNet (Hafner et al., 2019) and was employed for online planning inthe latent space.
Revised: RSSM was introduced in PlaNet (Hafner et al., 2019) and was employed for online planning in the latent space.

Original: DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021) achievedstate-of-the-art results on DMC (Tassa et al., 2018) and Atari (Bellemare et al., 2013), respectively.
Revised: DreamerV1 (Hafner et 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.

Revised: SimPLe (Kaiser et al., 2019) trains a PPO (Schulman et al., 2017) agent on the learned video generation model in pixel space.

Revised: SOLAR (Zhang et al., 2019) solved robotics tasks via guided policy search.

Original: Sparsity in RNN : Neural networks have widely adopted sparsity to reduce the memory footprintof weights and activations (LeCun et al., 1990; Chen et al., 2015; Han et al., 2015).
Revised: 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).

Original: Several workshave explored sparsity in RNNs.
Revised: Several works have explored sparsity in RNNs.

Original: Campos et al.
Revised: Campos et al.

Original: (2017) introduced a mechanism in RNNs that learnsto 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.
Revised: (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.

Original: Goyal et al.(2019) presented Recurrent Independent Mechanism (RIM), an architecture that consists of separaterecurrent modules which are sparsely updated using a learned attention mechanism.
Revised: Goyal et al. (2019) presented Recurrent Independent Mechanism (RIM), an architecture that consists of separate recurrent modules which are sparsely updated using a learned attention mechanism.

Original: In contrast to
Revised: In contrast to

Original: RIM, the number of updated state variables in VSG algorithm is not fixed.
Revised: RIM, the number of updated state variables in VSG algorithm is not fixed.

Original:  Discussion
Revised:  Discussion

Original: In this work, we introduce VSG and SVSG, two latent dynamics models leveraging sparse stateupdates.
Revised: In this work, we introduce VSG and SVSG, two latent dynamics models leveraging sparse state updates.

Original: The sparse update prior was found to facilitate more efficient behaviors in tasks requiringlong-horizon planning.
Revised: The sparse update prior was found to facilitate more efficient behaviors in tasks requiring long-horizon planning.

Original: Furthermore, SVSG is a purely stochastic model with a single componentin the model state.
Revised: Furthermore, SVSG is a purely stochastic model with a single component in the model state.

Original: We show that VSG and SVSG can outperform leading agents on the proposed
Revised: We show that VSG and SVSG can outperform leading agents on the proposed

Original: BringBackShapes task, a challenging partially-observable and stochastic environment.
Revised: BringBackShapes task, a challenging partially-observable and stochastic environment.

Original: BBS allowsfor controlling different factors of variation like stochasticity and partial-observability.
Revised: BBS allows for controlling different factors of variation like stochasticity and partial-observability.

Original: Experimentsconducted on various variations in BBS demonstrate that the proposed agents are more robust tonoise in the environment and can better retain information of seen objects.
Revised: Experiments conducted on various variations in BBS demonstrate that the proposed agents are more robust to noise in the environment and can better retain information of seen objects.

Original: Some limitations andpotential research directions for future research are as follows:
Revised: Some limitations and potential research directions for future research are as follows:

Original: • In the current implementation of VSG, the latent space does not exhibit disentanglement whichcould be an interesting direction for future research.
Revised: • In the current implementation of VSG, the latent space does not exhibit disentanglement which could be an interesting direction for future research.

Original: Gating mechanisms in VSG can also be combined with other recurrent architectures like RIM (Goyal et al., 2019).
Revised: Gating mechanisms in VSG can also be combined with other recurrent architectures like RIM (Goyal et al., 2019).

Original: • In this work, BBS was explored with only 2 factors of variation: partial-observability andstochasticity.
Revised: • In this work, BBS was explored with only 2 factors of variation: partial-observability and stochasticity.

Original: More controllable factors like the nature of entities (shape, size, color of objects),underlying physics (mass, friction, elasticity), or procedural background generation can beintroduced to further study generalization capabilities of RL agents.• SVSG being a purely stochastic model can further be used to estimate state uncertainty bymarginalizing over multiple samples paths to efficiently explore in an unknown environment.•
Revised: 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. • 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.

Revised: • Evaluation on first-person view 3D games like tasks in DMLab (Beattie et al., 2016) would be interesting.

Original: A 3D version of the BBS environment would be interesting as the viewpoints may also changewith rotation of agent and the underlying physics will make the task more challenging.•
Revised: 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.

Original:  We have used small latent dimensions and it would be interesting to train such models with largerarchitectures and on more complex tasks.
Revised: • We have used small latent dimensions and it would be interesting to train such models with larger architectures and on more complex tasks.

Original: Scaling the current architecture would also requireoptimizing the implementation to make them computationally feasible.• Categorical latents outperformed Gaussian latents as the stochastic states of RSSM (Hafner et al.,2021), especially for discrete control tasks.
Revised: Scaling the current architecture would also require optimizing the implementation to make them computationally feasible. • Categorical latents outperformed Gaussian latents as the stochastic states of RSSM (Hafner et al., 2021), especially for discrete control tasks.

Original: However, SVSG was not found to work well withCategorical latents and we believe that sampled sparse states are hard to optimize.•
Revised: However, SVSG was not found to work well with Categorical latents and we believe that sampled sparse states are hard to optimize.

Original:  Model-based RL for visual control is still in early stages.
Revised: • Model-based RL for visual control is still in early stages.

Original: However, a major challenge withdeploying such models in the real world is safety especially during exploration.
Revised: However, a major challenge with deploying such models in the real world is safety especially during exploration.

Original: This wouldrequire an accurate world model that allows learning policies with stringent safety constraintsthat avoid mistakes when deployed in the real world.
Revised: This would require an accurate world model that allows learning policies with stringent safety constraints that avoid mistakes when deployed in the real world.

Original: Such algorithms will rely on models that are robust when transferred from simulation to the real world.
Revised: Such algorithms will rely on models that are robust when transferred from simulation to the real world.

Original: Checklist
Revised: Acknowledgements

Original: The checklist follows the references. Please read the checklist guidelines carefully for information onhow to answer these questions. For each question, change the default [TODO] to [Yes] , [No] or[N/A] .
Revised: 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

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description. For example: • Did you include the license to the code and datasets? [Yes] See Section ?? .
Revised: Bhati and Vineet Jain.

Original: • Did you include the license to the code anddatasets?
Revised: The authors are also grateful to CIFAR for funding and

Original: [No] The code and the data are proprietary. • Did you include the license to the code and datasets? [N/A]
Revised: the Digital Research Alliance of Canada for computing resources.

Original: Please do not modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructionsblock and only keep the Checklistsection heading above along with the questions/answers below.
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] See Section 6.
Revised: [Yes] See Section 6.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] Describedin the section with limitations.
Revised: [Yes] Described in the section with limitations.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A] (b) Did you include complete proofs of all theoretical results?
Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

Original: [N/A] 
Revised: [N/A] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] The code for the model and dataset is in the supplementary material.
Revised: [Yes] The code for the model and dataset is in the supplementary material.

Original: There is a README file withthe instructions to run them.
Revised: There is a README file with the instructions to run them.

Original: (b) Did you specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [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 experiments multiple times)?
Revised: [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 experiments multiple times)?

Original: [Yes] All the results and plots presented in Section 4 and Appendix were obtained after training on multiple seeds ranging from 3-5.
Revised: [Yes] All the results and plots presented in Section 4 and Appendix were obtained after training on multiple seeds ranging from 3-5.

Original: (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)?
Revised: (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)?

Original: [Yes] We have mentioned about the GPUs and time taken to run on a single seed in the implementation details for each environment.
Revised: [Yes] We have mentioned about the GPUs and time taken to run on a single seed in the implementation details for each environment.

Original: 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?
Revised: 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?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [No] We used the implementation ofDreamerV2, DreamerPro and DrQ-v2 provided by the authors with MIT license.
Revised: [No] We used the implementation of DreamerV2, DreamerPro and DrQ-v2 provided by the authors with MIT license.

Original: For comparison with DBC, we thank the authors for sharing the evaluation logs.
Revised: For comparison with DBC, we thank the authors for sharing the evaluation logs.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes] The code for the new environment (BBS) is included in the supplementary material.
Revised: [Yes] The code for the new environment (BBS) is included in the supplementary material.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Original: [Yes] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: [Yes] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A] 
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]