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

Original: Abstract
Revised: Abstract

Original: Learning world models from their sensory inputs enables agents to plan for actions by imagining their future outcomes.
Revised: Learning world models from their sensory inputs enables agents to plan for actions by imagining their future outcomes.

Original: 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.
Revised: 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: In environments with many objects, often only a small number of them are moving or interacting at the same time.
Revised: In environments with many objects, often only a small number of them are moving or interacting at the same time.

Original: In this paper, we investigate integrating this inductive bias of sparse interactions into the latent dynamics of world models trained from pixels.
Revised: In this paper, we investigate integrating this inductive bias of sparse interactions into the latent dynamics of world models trained from pixels.

Original: First, we introduce Variational Sparse Gating (VSG), a latent dynamics model that updates its feature dimensions sparsely through stochastic binary gates.
Revised: First, we introduce Variational Sparse Gating (VSG), a latent dynamics model that updates its feature dimensions sparsely through stochastic binary gates.

Original: 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.
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: 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.
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 space without 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 sequences in the compact latent space 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 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).
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 recently proposed
Revised: Motivated by these benefits, the recently proposed

Original: DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021)
Revised: DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021)

Original: agents achieved state-of-the-art results on a wide range of visual control tasks.
Revised: 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 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.
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 Dreamer agents (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., 2019) 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; Bengio et 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 (Campos36th 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.
Revised: Different ways of applying sparse updates in RNNs have been investigated (Campos et al., 2017; Neil et al., 2016; Goyal et al., 2019), enabling a subset of state dimensions to be constant 36th Conference on Neural Information Processing Systems (NeurIPS 2022). during the update.

Original: 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.
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 a physical 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 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.
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 the model state was composed of two paths, an 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 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.
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: (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.
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 single component can achieve competitive results, and call it Simple
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

Original: Variational Sparse Gating (SVSG).
Revised: Variational Sparse Gating (SVSG).

Original: 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.
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 RL do 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 (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.
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 not allow 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 stochastic environment, called BringBackShapes (BBS), where the task is to push objects to a predefined goal area.
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 and avoid 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 and model-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 to demonstrate that the proposed agents have better memory for tracking observed objects and are more robust 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 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).
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 recurrent states 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 empirical evaluation 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 tasks when compared with agents that also use a deterministic component.
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.

Original: • BringBackShapes : We developed the BringBackShapes (BBS) environment to evaluate agents on partially-observable and stochastic settings where these variations can be controlled.
Revised: • BringBackShapes : We developed the BringBackShapes (BBS) environment to evaluate agents on partially-observable and stochastic settings where these variations can be controlled.

Original: Our experiments show that the proposed agents are more robust 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 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.
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 that maximizes the expected discounted sum of rewards E p ( (cid:80)
Revised: The goal is to learn a policy that maximizes the expected discounted sum of rewards E p ( (cid:80)

Original: Tt =1 γ t r t ) , where γ is the discount factor.
Revised: 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) encode a 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: The agents behavior (Appendix B) is derived 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:  𝑠𝑠𝑠 3 𝑠𝑎Ƹ𝑟𝑥𝑥𝑥ො𝑥𝑎𝑎Ƹ𝑟ො𝑥
Revised: While training, the world model is learned with collected experience, the policy is improved on 𝑠𝑠𝑠 3 𝑠𝑎Ƹ𝑟𝑥𝑥𝑥ො𝑥𝑎𝑎Ƹ𝑟ො𝑥

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

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

Original: VSG VSG VSG Ƹ𝑟 3 ො𝑎ො𝑎ො𝑎ො𝑣ො𝑣ො𝑣Ƹ𝑠Ƹ𝑠(b)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.
Revised: VSG VSG VSG Ƹ𝑟 3 ො𝑎ො𝑎ො𝑎ො𝑣ො𝑣ො𝑣Ƹ𝑠Ƹ𝑠(b) 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 training progresses, 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 collected experience and 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 the world and an action, the model applies the learned transition dynamics to predict the resulting next state 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 was composed 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 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).
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 subset of 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 enable the agent to have long-term memory and learn robust representations to solve 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 predictors for 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 t using Convolutional Neural Networks (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 model equipped with the stochastic gating mechanism to get the recurrent state h t , 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 posterior representation state z t is obtained using the representation model and contains information about the current observation x t .
Revised: The posterior representation state z t is obtained using the representation model and contains information about the current observation x t .

Original: The prior state ˆ z t is obtained from the transition predictor without observing the current observation x t .
Revised: The prior state ˆ z t is obtained from the transition predictor without observing the current observation x t .

Original: This is useful while planning as sequences are generated in compact latent state, and the output from the transition predictor is utilized.
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 also results in a lower memory footprint and enables predictions of thousands 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 a known distribution with learned parameters like Gaussian (Hafner et al., 2020) or Categorical (Hafner et 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 gives the 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 the original 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 predict the probability that an episode will end.
Revised: The discount factor helps to predict the probability that an episode will end.

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

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

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

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

Original: 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 .
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: The transition predictor provides the prior image representation state ˆ z t .
Revised: The transition predictor provides the prior image representation state ˆ z t .

Original: 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.
Revised: 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 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.
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 the components 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 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.
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 were found to alleviate the vanishing gradient problem by effectively reducing the number of sequential operations (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 avoiding the gradual change of state values introduced by repeated multiplication with continuous gate values in 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 take binary 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 through 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 is a reset and update gate.
Revised: Similar to GRU (Cho et al., 2014), there is a reset and update gate.

Original: 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.
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 only binary 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: 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 .
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 (Bengio et 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:
Revised: The update equations are:

Original: [Equation]
Revised: [Equation]

Original: 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.
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 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.
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 unit variance, 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 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.
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 (Hafner et 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 loss to 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 world model 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]
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 was discussed 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 information for 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 this work, 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: 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. 

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

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

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

Original: We have used a Gaussian distribution for the stochastic state with a learnable mean vector and a learnable diagonal covariance matrix.
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 reward and 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 model in SVSG are:
Revised: The components of world model in SVSG are:

Original: Representation model:
Revised: Representation model:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Reward predictor
Revised: Reward predictor

Original: [Equation]
Revised: [Equation]

Original: The representation model q ϕ and transition predictor p ϕ are modified to output the posterior s t and prior ˆ 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.

Original: 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 .
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 t and 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 distribution to sparsely update the latent states at each step, given by:
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 tangent activation 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 is feeded 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 was predicted without it.
Revised: The image encoding x t was used to get posterior distribution, whereas the prior distribution was predicted without it.

Original: 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.
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 and posterior, respectively.
Revised: where f p and f q denotes functions that output a distribution with learnable parameters for prior and posterior, respectively.

Original: For SVSG, Categorical latents (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 that samples from a categorical distribution are binary vectors and it is difficult to accurately reconstruct with 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 that sparse gating mechanism introduced in VSG was important for convergence of SVSG.
Revised: Lastly, we observed that sparse gating mechanism introduced in VSG was important for convergence of SVSG.

Original: 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).
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 corresponding time 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 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.
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:  Experiments 4.1 BringBackShapes
Revised:  Experiments 4.1 BringBackShapes

Original: 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.
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 push the 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 an obfuscated 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 efficiently explore 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 account for 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
Revised: Stochasticity was introduced in the environment by

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

Revised: [Figure]

Original: [Table]
Revised: [Table]

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

Original: The distractor objects follow Brownian motion in any direction and can’t be 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 push objects 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: Due to partial-observability, such instances might not be visible to the agent, making the task even more challenging.
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: Experimental Setup : The BringBackShapes (BBS) environment returns high dimensional images of shape 64 × 64 × 3 as observation.
Revised: Experimental Setup : The BringBackShapes (BBS) environment returns high dimensional images of shape 64 × 64 × 3 as observation.

Original: Action is a 2-dimensional continuous vector with acceleration and direction as components.
Revised: Action is a 2-dimensional continuous vector with acceleration and direction as components.

Original: Episodes last for 3000 environment steps and an action repeat (Mnih et al., 2016) of 4 was used.
Revised: Episodes last for 3000 environment steps and an action repeat (Mnih et al., 2016) of 4 was used.

Original: Baseline agents include DreamerV2 (Hafner et al., 2021) and DrQ-v2 (Yarats et al., 2022).
Revised: Baseline agents include DreamerV2 (Hafner et al., 2021) and DrQ-v2 (Yarats et al., 2022).

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

Original: SVSG.
Revised: SVSG.

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

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

Original: Lastly, results are reported across 5 seeds.
Revised: Lastly, results are reported across 5 seeds.

Original: Quantitative Results : Fig.
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 indicates how 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 3, we trained
Revised: As discussed in Section 3, we trained

Original: SVSG with Gaussian latents only.
Revised: SVSG with Gaussian latents only.

Original: Upon evaluation at 2.5M timesteps, DreamerV2 achieves 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, DreamerV2 does 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 with sparsity 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 DreamerV but similar to VSG, demonstrating that a purely stochastic model can achieve similar performance.
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 observability and 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 the portion 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 number  of distractor objects.
Revised: 3 [c] presents the plots of models trained at 2.5M steps at different sizes of the arena and number [Figure] of distractor objects.

Original: It can be observed that increasing the size of arena makes it harder to score objects.
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 not[Figure] perform well on larger arena sizes.
Revised: However, SVSG did not perform well on larger arena sizes.

Original: Adding noisy distractor objects led to drop in final performance of 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 prior helps 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 completely visible in the agent’s view) and is calculated by averaging the first-visit time of each object in the arena.
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 indicates that an agent is able to quickly discover 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 of steps 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: 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.
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 of visited objects that were not scored which indicates us that the agent might not be remembering positions of 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 to an 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 the environment 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 to the 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 did not follow a smooth trajectory and was spending more time around an object to push it in the goal.
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, our methods showed smoother trajectories and were more efficient 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, 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
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: Fig. 4 and supplementary material for more videos).
Revised: Fig. 4 and supplementary material for more videos).

Original: Effect of Sparse Gating :
Revised: Effect of Sparse Gating :

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

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

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

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

Original: SVSG remembers the color and location of objects, and is also cognizant about the goal location and walls.
Revised: SVSG remembers the color and location of objects, and is also cognizant about the goal location and walls.

Original: Furthermore, unrolled trajectories from the world model of DreamerV2 showed distortion in the shapes, and in some instances modifies the color of the objects.
Revised: Furthermore, unrolled trajectories from the world model of DreamerV2 showed distortion in the shapes, and in some instances modifies the color of the objects.

Original: 4.2 DeepMind Control Suite
Revised: 4.2 DeepMind Control Suite

Original: Experimental Setup :
Revised: Experimental Setup :

Original: The proposed method is evaluated on a few tasks from DeepMind Control
Revised: The proposed method is evaluated on a few tasks from DeepMind Control

Original: Suite (Tassa et al., 2018).
Revised: Suite (Tassa et al., 2018).

Original: 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.
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 was used.
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 trained on 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 for 1M environment steps.
Revised: The agents were trained for 1M environment steps.

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

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 trained using 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 the respective 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: Results : Figure 5 presents comparison of VSG and SVSG with the baseline agents (Refer to Appendix D for comparison on 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 performs comparable 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, 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.
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, 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: Ablation Studies : 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 with other forms of noise where natural videos are used in the background for
Revised: We also experimented with other forms of noise where natural videos are used in the background for

Original: DeepMind Control tasks (Zhang et al., 2021; 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) 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
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

Original: Appendix F).
Revised: Appendix F).

Original: We also experimented with VSG in discrete control tasks from the Atari benchmark in
Revised: We also experimented with VSG in discrete control tasks from the Atari benchmark in

Original: Appendix G, where VSG performed better on task with changing viewpoints.
Revised: Appendix G, where VSG performed better on task with changing viewpoints.

Original:  Related Work
Revised:  Related Work

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

Original: 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.
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: 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).
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 StateSpace Model (RSSM) (Hafner et al., 2019) comprises 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
Revised: Hierarchical latent models for video prediction were proposed in

Original: CWVAEs (Saxena et al., 2021) with levels ticking at different intervals in time.
Revised: CWVAEs (Saxena et al., 2021) with levels ticking at different intervals in time.

Original: (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.
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 (Ha and Schmidhuber, 2018; Hafner et al., 2019; 2020; 2021;
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;

Original: 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;
Revised: 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: Schwarzer et al., 2021;
Revised: Schwarzer et al., 2021;

Original: Mondal et al., 2022).
Revised: Mondal et al., 2022).

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

Original: TD3 (Fujimoto et al., 2018) builds upon the DDPG algorithm and addresses the problem of overestimation bias in the value function.
Revised: 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 contrastive losses 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 (Yarats et 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: RSSM was introduced in PlaNet (Hafner et al., 2019) and was employed for online planning in the 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) achieved state-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.

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

Original: SOLAR (Zhang et al., 2019) solved robotics tasks via guided policy search.
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 footprint of 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 works have 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 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.
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.
Revised: Goyal et al.

Original: (2019) presented Recurrent Independent Mechanism (RIM), an architecture that consists of separate recurrent modules which are sparsely updated using a learned attention mechanism.
Revised: (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 state updates.
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 requiring long-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 component in 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 allows for 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: 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.
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 and potential 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 which could 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 and stochasticity.
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 be introduced to further study generalization capabilities of RL agents.
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.

Original: • 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: • 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.

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

Original: 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.
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 larger architectures 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 require optimizing the implementation to make them computationally feasible.
Revised: Scaling the current architecture would also require optimizing the implementation to make them computationally feasible.

Original: • Categorical latents outperformed Gaussian latents as the stochastic states of RSSM (Hafner et al., 2021), especially for discrete control tasks.
Revised: • 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 with Categorical 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 with deploying 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 would require an accurate world model that allows learning policies with stringent safety constraints that 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: Acknowledgements
Revised: Acknowledgements

Original: The authors would like to thank David Meger, Lucas Lehnert and Ankesh Anand for their valuable feedback and discussions.
Revised: The authors would like to thank David Meger, Lucas Lehnert and Ankesh Anand for their valuable feedback and discussions.

Original: The text also benefited from discussions with Abhinav Agarwalla, Rupali
Revised: The text also benefited from discussions with Abhinav Agarwalla, Rupali

Original: Bhati and Vineet Jain.
Revised: Bhati and Vineet Jain.

Original: The authors are also grateful to CIFAR for funding and
Revised: The authors are also grateful to CIFAR for funding and

Original: the Digital Research
Revised: the Digital Research

Original: Alliance of Canada for computing resources.
Revised: Alliance of Canada for computing resources

Original: Checklist
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] Described in 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] 2.
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] 3.
Revised: [N/A] 3.

Original: 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)?
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 with the 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 of DreamerV2, 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] 5.
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, if applicable?
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]