Original: F ACTORIZED W ORLD M ODELS FOR L EARNING C AUSAL R ELATIONSHIPS
Revised: F ACTORIZED W ORLD M ODELS FOR L EARNING C AUSAL R ELATIONSHIPS

Original: A BSTRACT
Revised: A BSTRACT

Original: World models serve as a powerful framework for model-based reinforcement learning, and they can greatly benefit from the shared structure of the world environments.
Revised: World models serve as a powerful framework for model-based reinforcement learning, and they can greatly benefit from the shared structure of the world environments.

Original: However, learning the high-level causal influence of objects on each other remains a challenge.
Revised: However, learning the high-level causal influence of objects on each other remains a challenge.

Original: In this work, we propose CEMA, a structured world model with factorized latent state capable of modeling sparse interaction.
Revised: In this work, we propose CEMA, a structured world model with factorized latent state capable of modeling sparse interaction.

Original: This is possible due to a separate state and dynamics of three components: the actor, the object of manipulation, the latent influence factor between these two states.
Revised: This is possible due to a separate state and dynamics of three components: the actor, the object of manipulation, the latent influence factor between these two states.

Original: In multitask setting, we analyze the mutual information of the hierarchical latent states to show how the model can represent sparse updates and directly model the causal influence of the robot on the object
Revised: In multitask setting, we analyze the mutual information of the hierarchical latent states to show how the model can represent sparse updates and directly model the causal influence of the robot on the object

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Reinforcement learning (RL) has recently benefited extensively from structural inductive biases.
Revised: Reinforcement learning (RL) has recently benefited extensively from structural inductive biases.

Original: Assuming block structure of MDP (Du et al., 2021) or graph-structured objects interaction model (Veerapaneni et al., 2019) one can improve generalization in RL (Kirk et al., 2021).
Revised: Assuming block structure of MDP (Du et al., 2021) or graph-structured objects interaction model (Veerapaneni et al., 2019) one can improve generalization in RL (Kirk et al., 2021).

Original: Other works bring structure into the world model of the agent to facilitate generalization, sample-efficiency, or casual modeling (Wang et al., 2020; Hallak et al., 2015).
Revised: Other works bring structure into the world model of the agent to facilitate generalization, sample-efficiency, or casual modeling (Wang et al., 2020; Hallak et al., 2015).

Original: The latter is especially promising as world models (Ha & Schmidhuber, 2018) can tremendously benefit from the reuse of the learned or encoded structure of the world.
Revised: The latter is especially promising as world models (Ha & Schmidhuber, 2018) can tremendously benefit from the reuse of the learned or encoded structure of the world.

Original: In this work, we present CEMA (Cause-Effect Modeling Agent), a world model that learns the correct cause-effect relationship between the actor and the object of manipulation.
Revised: In this work, we present CEMA (Cause-Effect Modeling Agent), a world model that learns the correct cause-effect relationship between the actor and the object of manipulation.

Original: We show how the model can learn latent influence from the actor state on the object state.
Revised: We show how the model can learn latent influence from the actor state on the object state.

Original: We hypothesize that the state of the manipulator should be modeled separately from the state of the object of manipulation.
Revised: We hypothesize that the state of the manipulator should be modeled separately from the state of the object of manipulation.

Original: Each of them should have separate latent dynamics where actions control the former, and the latter is controlled by an abstract influence factor determining how the object state should be updated in response to the manipulator state update.
Revised: Each of them should have separate latent dynamics where actions control the former, and the latter is controlled by an abstract influence factor determining how the object state should be updated in response to the manipulator state update.

Original: This latent factor can be seen as a high-level action that directly controls the object state, the idea similar to (Eysenbach et al., 2021).
Revised: This latent factor can be seen as a high-level action that directly controls the object state, the idea similar to (Eysenbach et al., 2021).

Original: We build the world model by employing these latent variables in a hierarchical fashion, and we show that prior distribution models the forward cause-effect relationship while the approximate posterior models the inverse process of reducing the effect to cause.
Revised: We build the world model by employing these latent variables in a hierarchical fashion, and we show that prior distribution models the forward cause-effect relationship while the approximate posterior models the inverse process of reducing the effect to cause.

Original: In our experiments, we fix the setup to a known common cause case, so we do not need to discover it as it is fixed and known.
Revised: In our experiments, we fix the setup to a known common cause case, so we do not need to discover it as it is fixed and known.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Model-based RL algorithms can significantly benefit from having a structured world model, especially in a visual control setting.
Revised: Model-based RL algorithms can significantly benefit from having a structured world model, especially in a visual control setting.

Original: PlaNet (Hafner et al., 2018) and Dreamer (Hafner et al., 2019) train a recurrent world model, separating deterministic and stochastic parts of environment dynamics.
Revised: PlaNet (Hafner et al., 2018) and Dreamer (Hafner et al., 2019) train a recurrent world model, separating deterministic and stochastic parts of environment dynamics.

Original: However, RPC (Eysenbach et al., 2021) argues that robust and generalizable behaviors also depend on compression in representation learning.
Revised: However, RPC (Eysenbach et al., 2021) argues that robust and generalizable behaviors also depend on compression in representation learning.

Original: Following the Information Bottleneck principle (Tishby et al., 2000;
Revised: Following the Information Bottleneck principle (Tishby et al., 2000;

Original: Alemi et al.,
Revised: Alemi et al.,

Original: , they build a compressed representation and restrict a policy to follow predictable paths.
Revised: , they build a compressed representation and restrict a policy to follow predictable paths.

Original: We follow a similar principle, combining it with a causal-based, object-centric world model.
Revised: We follow a similar principle, combining it with a causal-based, object-centric world model.

Original: The compound causal structure is necessary for modeling good object interactions.
Revised: The compound causal structure is necessary for modeling good object interactions.

Original: C-SWM (Kipf et al., 2020) adopts graph neural network for object-level reasoning, allowing for flexible relationship modeling.
Revised: C-SWM (Kipf et al., 2020) adopts graph neural network for object-level reasoning, allowing for flexible relationship modeling.

Original: However, there is no explicitly stated causality between objects.
Revised: However, there is no explicitly stated causality between objects.

Original: OP3 (Veerapaneni et al., 2019) operates on entities rather than objects to simulate laws of the environment, equally valid for every object in it.
Revised: OP3 (Veerapaneni et al., 2019) operates on entities rather than objects to simulate laws of the environment, equally valid for every object in it.

Original: While they achieve impressive generalization in combinatorically hard problems, modeling asymmetric or sparse causal relationships is not possible.
Revised: While they achieve impressive generalization in combinatorically hard problems, modeling asymmetric or sparse causal relationships is not possible.

Original: In GATSBI (Min et al., 2021) authors model agent, background, and objects differently.
Revised: In GATSBI (Min et al., 2021) authors model agent, background, and objects differently.

Original: The influence of entities on each other is designed as GNN, taking into account their different nature.
Revised: The influence of entities on each other is designed as GNN, taking into account their different nature.

Original: Whereas this approach is suitable for scenes with numerous simple objects, it may lack the expressiveness for the case of complex actor-object interaction.
Revised: Whereas this approach is suitable for scenes with numerous simple objects, it may lack the expressiveness for the case of complex actor-object interaction.

Original: A number of works on representation learning in object-centric context (Greff et al., 2019; Locatello et al., 2020; Engelcke et al., 2020) also use hierarchical latents to introduce a structure into the scene.
Revised: A number of works on representation learning in object-centric context (Greff et al., 2019; Locatello et al., 2020; Engelcke et al., 2020) also use hierarchical latents to introduce a structure into the scene.

Original: However, they do not explore modeling dynamics in these terms.
Revised: However, they do not explore modeling dynamics in these terms.

Original:  B ACKGROUND 3. R EINFORCEMENT L EARNING
Revised:  B ACKGROUND 3. R EINFORCEMENT L EARNING

Original: [Figure]
Revised: [Figure]

Original: We formalize the visual control problem as a partially observable
Revised: We formalize the visual control problem as a partially observable

Original: Markov Decision Process M = ⟨S , A , O , P, R, γ ⟩ , where S is the set of states, A is the set of actions, and O is the set of observations.
Revised: Markov Decision Process M = ⟨S , A , O , P, R, γ ⟩ , where S is the set of states, A is the set of actions, and O is the set of observations.

Original: The state dynamics follow the conditional transition distribution P ( s ′ | s, a ) , but the agent only has access to samples from the observation distribution o ∼ p ( o | s ) , o ∈ O .
Revised: The state dynamics follow the conditional transition distribution P ( s ′ | s, a ) , but the agent only has access to samples from the observation distribution o ∼ p ( o | s ) , o ∈ O .

Original: The agent is defined as a policy π ( a t | o ≤ t , a <t ) .
Revised: The agent is defined as a policy π ( a t | o ≤ t , a <t ) .

Original: It interacts with a partially observable environment by taking actions on the environment and getting a next observation.
Revised: It interacts with a partially observable environment by taking actions on the environment and getting a next observation.

Original: We write o t , r t ∼ p ( o t , r t | o <t , a <t ) as a shorthand for s t ∼ P ( s t | s t − 1 , a t − 1 ) , o t =
Revised: We write o t , r t ∼ p ( o t , r t | o <t , a <t ) as a shorthand for s t ∼ P ( s t | s t − 1 , a t − 1 ) , o t =

Original: O ( s t , a t − 1 ) , r t =
Revised: O ( s t , a t − 1 ) , r t =

Original: R ( s t , a t − 1 ) .
Revised: R ( s t , a t − 1 ) .

Original: The goal of the agent is to maximize its expected discounted sum of rewards R ( π, M ) =
Revised: The goal of the agent is to maximize its expected discounted sum of rewards R ( π, M ) =

Original: E π (cid:80) t γ t r t .
Revised: E π (cid:80) t γ t r t .

Original: We also adopt notion of Context MDP (CMDP) introduced in (Kirk et al., 2021).
Revised: We also adopt notion of Context MDP (CMDP) introduced in (Kirk et al., 2021).

Original: A context is a latent variable that determines the task in CMDP.
Revised: A context is a latent variable that determines the task in CMDP.

Original: Each context defines the states that could be visited during an episode such as a level in level-based environments or the parametric environment configuration in environments that are parameterized by vector.
Revised: Each context defines the states that could be visited during an episode such as a level in level-based environments or the parametric environment configuration in environments that are parameterized by vector.

Original: Context variable allows to factorize the initial state distribution p ( s )
Revised: Context variable allows to factorize the initial state distribution p ( s )

Original: =
Revised: =

Original: p ( s ′ | c ) p ( c ) , where s = ( s ′ , c ) .
Revised: p ( s ′ | c ) p ( c ) , where s = ( s ′ , c ) .

Original: The observation is then sampled from the emission distribution, i.e., o ∼ p ( o | s ′ ) .
Revised: The observation is then sampled from the emission distribution, i.e., o ∼ p ( o | s ′ ) .

Original: 3. M ODEL -B ASED R EINFORCEMENT L EARNING VIA W ORLD M ODELS
Revised: 3. M ODEL -B ASED R EINFORCEMENT L EARNING VIA W ORLD M ODELS

Original: World models explicitly learn environment dynamics to generate novel virtual experience (Sutton, 1991;
Revised: World models explicitly learn environment dynamics to generate novel virtual experience (Sutton, 1991;

Original: Ha & Schmidhuber, 2018).
Revised: Ha & Schmidhuber, 2018).

Original: When observations are non-Markovian, e.g., in visual control, this can be achieved only by incorporating latent states.
Revised: When observations are non-Markovian, e.g., in visual control, this can be achieved only by incorporating latent states.

Original: For latent dynamics learning, we use the RSSM (Hafner et al., 2018), which learns the dynamics by building a Markovian latent state for each timestep s t given previous action a t − 1 by autoencoding observations o t and rewards r t , which are non-Markovian.
Revised: For latent dynamics learning, we use the RSSM (Hafner et al., 2018), which learns the dynamics by building a Markovian latent state for each timestep s t given previous action a t − 1 by autoencoding observations o t and rewards r t , which are non-Markovian.

Original: Dreamer’s world model consists of a representation model, or an encoder q ( s t | s t − 1 , a t − 1 , o t ) , a transition model p ( s t | s t − 1 , a t − 1 ) , an observation model p ( o t
Revised: Dreamer’s world model consists of a representation model, or an encoder q ( s t | s t − 1 , a t − 1 , o t ) , a transition model p ( s t | s t − 1 , a t − 1 ) , an observation model p ( o t

Original: | s t ) , and a reward model p ( r t | s t ) .
Revised: | s t ) , and a reward model p ( r t | s t ) .

Original: The representation and transition models share common parameters in the RSSM network (Hafner et al., 2018).
Revised: The representation and transition models share common parameters in the RSSM network (Hafner et al., 2018).

Original: The world model is trained to maximize the variational lower bound on the likelihood (ELBO) of the observed trajectory conditioned on the actions E p log p ( o 1: T , r 1: T | a 1: T ) .
Revised: The world model is trained to maximize the variational lower bound on the likelihood (ELBO) of the observed trajectory conditioned on the actions E p log p ( o 1: T , r 1: T | a 1: T ) .

Original: This is done by incorporating an approximate posterior model q ( s t | s t − 1 , a t − 1 , o t ) , which is known as a representation model.
Revised: This is done by incorporating an approximate posterior model q ( s t | s t − 1 , a t − 1 , o t ) , which is known as a representation model.

Original: This model acts as a proposal distribution for states s t .
Revised: This model acts as a proposal distribution for states s t .

Original:  L ATENT STATE FACTORIZATION
Revised:  L ATENT STATE FACTORIZATION

Original: In this section, we describe a latent factorization we use in our method.
Revised: In this section, we describe a latent factorization we use in our method.

Original: In many cases, the scene components of the visual control problem can be divided into two groups based on their functional role: acting and objective .
Revised: In many cases, the scene components of the visual control problem can be divided into two groups based on their functional role: acting and objective .

Original: The first one represents everything that can be controlled directly by an agent, e.g., a robotic arm, as in our case.
Revised: The first one represents everything that can be controlled directly by an agent, e.g., a robotic arm, as in our case.

Original: The latter stands for everything that controls what the agent can interact with.
Revised: The latter stands for everything that controls what the agent can interact with.

Original: We model each group using hierarchical latent variables with an additional latent influence factor.
Revised: We model each group using hierarchical latent variables with an additional latent influence factor.

Original: To provide correct signal for modeling both parts, the input image
Revised: To provide correct signal for modeling both parts, the input image

Original: o t is separated into corresponding parts via a segmentation mask: the robot image o t and the object image o ′ t .
Revised: o t is separated into corresponding parts via a segmentation mask: the robot image o t and the object image o ′ t .

Original: The generative part of the model explicitly expresses the direct cause-effect relationships between the actor and the object, resulting in a bottom-up imagination process.
Revised: The generative part of the model explicitly expresses the direct cause-effect relationships between the actor and the object, resulting in a bottom-up imagination process.

Original: Since the actor is assumed to be directly controlled by the MDP dynamics, we model its updates using robot state s t with latent transition model p ( s t | s t − 1 , a t − 1 ) and observation model p ( o t
Revised: Since the actor is assumed to be directly controlled by the MDP dynamics, we model its updates using robot state s t with latent transition model p ( s t | s t − 1 , a t − 1 ) and observation model p ( o t

Original: | s t ) .
Revised: | s t ) .

Original: The objective latent states s ′ t , on the other hand, are not controlled directly via actions a t .
Revised: The objective latent states s ′ t , on the other hand, are not controlled directly via actions a t .

Original: Instead, their control is mediated by the influence of the physical parameters of the actor on the object at each point in time.
Revised: Instead, their control is mediated by the influence of the physical parameters of the actor on the object at each point in time.

Original: We express this collateral control using a latent influence factor u t that abstracts how the objective state should be changed to reflect the observational update.
Revised: We express this collateral control using a latent influence factor u t that abstracts how the objective state should be changed to reflect the observational update.

Original: Since the actor state does not contain the information about the current task, we condition the influence factor on the MDP context vector c in addition to the robot state s t , which results in an influence model of the form p ( u t | s t , c ) .
Revised: Since the actor state does not contain the information about the current task, we condition the influence factor on the MDP context vector c in addition to the robot state s t , which results in an influence model of the form p ( u t | s t , c ) .

Original: For the objective part, we use the prior distribution in the form of a Hidden Markov model (HMM) p ( s ′ t | s ′ t − 1 , u t ) to model influenced objective updates and the objective observation model p ( o ′ t | s ′ t ) .
Revised: For the objective part, we use the prior distribution in the form of a Hidden Markov model (HMM) p ( s ′ t | s ′ t − 1 , u t ) to model influenced objective updates and the objective observation model p ( o ′ t | s ′ t ) .

Original: The posterior model inverts causal relationship and infers states in a top-down fashion.
Revised: The posterior model inverts causal relationship and infers states in a top-down fashion.

Original: This hierarchical structure is motivated by the causal effects that the agent should model.
Revised: This hierarchical structure is motivated by the causal effects that the agent should model.

Original: We hypothesize that such architecture would allow the model to propagate the information from the cause, i.e., the robot, to an effect, i.e., the object movement.
Revised: We hypothesize that such architecture would allow the model to propagate the information from the cause, i.e., the robot, to an effect, i.e., the object movement.

Original: The object posterior model q ( s ′ t | o ′ t ) infers the actual updated state solely from the updated observation.
Revised: The object posterior model q ( s ′ t | o ′ t ) infers the actual updated state solely from the updated observation.

Original: This choice was made to account for sparse objective dynamics as the recurrent stochastic model would pass the previous state as is and not account for sparsity.
Revised: This choice was made to account for sparse objective dynamics as the recurrent stochastic model would pass the previous state as is and not account for sparsity.

Original: The inferred objective state is then used to produce the influence factor u t via the inverse influence distribution q ( u t | s ′ t ) .
Revised: The inferred objective state is then used to produce the influence factor u t via the inverse influence distribution q ( u t | s ′ t ) .

Original: This distribution reverses the causal process of influencing s t on s ′ t and produces a sample of u t that could lead to the actual updated object state s ′ t .
Revised: This distribution reverses the causal process of influencing s t on s ′ t and produces a sample of u t that could lead to the actual updated object state s ′ t .

Original: Using the embedding of the robot observation o t , the model infers an updated robot state s t given the inverse influence factor u t and the previous state-action pair s t − 1 , a t − 1 resulting in the robot posterior q ( s t | s t
Revised: Using the embedding of the robot observation o t , the model infers an updated robot state s t given the inverse influence factor u t and the previous state-action pair s t − 1 , a t − 1 resulting in the robot posterior q ( s t | s t

Original: [Equation]
Revised: [Equation]

Original: 1 , a t − 1 , u t , o t ) .
Revised: 1 , a t − 1 , u t , o t ) .

Original: The world model optimizes the following lower bound:
Revised: The world model optimizes the following lower bound:

Original: [Equation]
Revised: [Equation]

Original: We implement the acting prior and posterior dynamics via the RSSM model (Hafner et al., 2018).
Revised: We implement the acting prior and posterior dynamics via the RSSM model (Hafner et al., 2018).

Original: For the objective part, instead, we implement it using the following architecture:
Revised: For the objective part, instead, we implement it using the following architecture:

Original: [Equation]
Revised: [Equation]

Original: Here, e t +1 is a CNN embedding of the observation, and Dist is any distribution such as Normal, Categorical, or Bernoulli.
Revised: Here, e t +1 is a CNN embedding of the observation, and Dist is any distribution such as Normal, Categorical, or Bernoulli.

Original: The nature of the process dictates the structure of the objective models we represent with these distributions.
Revised: The nature of the process dictates the structure of the objective models we represent with these distributions.

Original: RSSM is an appropriate choice for the controlled processes where there is a direct causal factor, such as a t −for the actor modeling.
Revised: RSSM is an appropriate choice for the controlled processes where there is a direct causal factor, such as a t −for the actor modeling.

Original: Object state, however, could be changed quite rarely and is instead updated by aftereffects of a t − 1 .
Revised: Object state, however, could be changed quite rarely and is instead updated by aftereffects of a t − 1 .

Original: Therefore, we infer the latent object state only from its current observation o ′ t , while the influence factor u t acts as a proxy of actor history that led to s ′ t .
Revised: Therefore, we infer the latent object state only from its current observation o ′ t , while the influence factor u t acts as a proxy of actor history that led to s ′ t .

Original: The objective model takes form similar to the one in the RPC agent (Eysenbach et al., 2021).
Revised: The objective model takes form similar to the one in the RPC agent (Eysenbach et al., 2021).

Original: Note also, that an inverse influence factor affects the inferred robot dynamics.
Revised: Note also, that an inverse influence factor affects the inferred robot dynamics.

Original: This factor explains derailments of robot state caused by, e.g., bumping object.
Revised: This factor explains derailments of robot state caused by, e.g., bumping object.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Rotated Drawer World.
Revised: Rotated Drawer World.

Original: To test how the model can infer the correct causal relationship between the robot and the object of control, we built a custom robotic visual control environment based on MuJoCo (Todorov et al., 2012) and MetaWorld (Yu et al., 2019).
Revised: To test how the model can infer the correct causal relationship between the robot and the object of control, we built a custom robotic visual control environment based on MuJoCo (Todorov et al., 2012) and MetaWorld (Yu et al., 2019).

Original: Our environment represents a Sawyer robot that stands before the table with a green drawer.
Revised: Our environment represents a Sawyer robot that stands before the table with a green drawer.

Original: The robot’s goal is to open the drawer while the position of the drawer can vary between tasks.
Revised: The robot’s goal is to open the drawer while the position of the drawer can vary between tasks.

Original: We parameterize the robot’s task by a polar coordinate of the drawer.
Revised: We parameterize the robot’s task by a polar coordinate of the drawer.

Original: Namely, we fix the radius to make the task effectively dependent only on a polar angle, making the context set C = [0 , 360) ⊂ R .
Revised: Namely, we fix the radius to make the task effectively dependent only on a polar angle, making the context set C = [0 , 360) ⊂ R .

Original: Examples of the image observations are in Figure 2.
Revised: Examples of the image observations are in Figure 2.

Original: A more specific description of the environment is provided in Appendix A.
Revised: A more specific description of the environment is provided in Appendix A.

Original: The motivation to use this setup is as follows.
Revised: The motivation to use this setup is as follows.

Original: First, each task alone is simple for a visual RL algorithm to solve quickly.
Revised: First, each task alone is simple for a visual RL algorithm to solve quickly.

Original: Second, the tasks share much common structure with only one variation parameter.
Revised: Second, the tasks share much common structure with only one variation parameter.

Original: Third, the topdown camera allows observing the drawer from the same angle for each task leaving the need of understanding complicated 3D object variations.
Revised: Third, the topdown camera allows observing the drawer from the same angle for each task leaving the need of understanding complicated 3D object variations.

Original: These facts combined lead to the testbed that promotes fast learning of similar tasks and that is aimed at the reuse and understanding of the information extracted from those tasks.
Revised: These facts combined lead to the testbed that promotes fast learning of similar tasks and that is aimed at the reuse and understanding of the information extracted from those tasks.

Original: Baselines.
Revised: Baselines.

Original: As a baseline, we train the Dreamer agent on the proposed set of tasks.
Revised: As a baseline, we train the Dreamer agent on the proposed set of tasks.

Original: We leave all hyperparameters unchanged compared to the original algorithm (Hafner et al., 2019).
Revised: We leave all hyperparameters unchanged compared to the original algorithm (Hafner et al., 2019).

Original: The model is trained for a million steps on tasks with contexts sampled from the uniform distribution over the training contexts set.
Revised: The model is trained for a million steps on tasks with contexts sampled from the uniform distribution over the training contexts set.

Original: Results.
Revised: Results.

Original: The CEMA agent shows comparable performance with Dreamer in terms of final reward.
Revised: The CEMA agent shows comparable performance with Dreamer in terms of final reward.

Original: Both models converge to approximately the same value of the episode return ( > 3500 ).
Revised: Both models converge to approximately the same value of the episode return ( > 3500 ).

Original: The convergence plot is present in Appendix B.
Revised: The convergence plot is present in Appendix B.

Original: The return above 2000 corresponds to a robot correctly approaching the drawer; the 3000 value of the return corresponds to a correct opening of the drawer.
Revised: The return above 2000 corresponds to a robot correctly approaching the drawer; the 3000 value of the return corresponds to a correct opening of the drawer.

Original: The difference between the returns above 3000 is the speed at which the robot opens the drawer.
Revised: The difference between the returns above 3000 is the speed at which the robot opens the drawer.

Original: We found the non-recurrent object layer posterior q ( s ′ t | o ′ t ) to be performing much better compared to a recurrent variant.
Revised: We found the non-recurrent object layer posterior q ( s ′ t | o ′ t ) to be performing much better compared to a recurrent variant.

Original: This is probably due to the fact that the object updates are sparse and are performed for a short period of time under certain conditions (when the robot interacts with the object).
Revised: This is probably due to the fact that the object updates are sparse and are performed for a short period of time under certain conditions (when the robot interacts with the object).

Original: Also, the inverse cause nature of the posterior distribution showed to be contributing to the final performance.
Revised: Also, the inverse cause nature of the posterior distribution showed to be contributing to the final performance.

Original: As the main contribution in this work, we show that the model reflects cause-effect relationship between the robot and the object of control.
Revised: As the main contribution in this work, we show that the model reflects cause-effect relationship between the robot and the object of control.

Original: We found that this fact can qualitatively be expressed by studying mutual information (MI) between the input and output of each distribution.
Revised: We found that this fact can qualitatively be expressed by studying mutual information (MI) between the input and output of each distribution.

Original: For example, for the influence model prior, we calculate  I ( u t ,( s t, c ) | t ) =
Revised: For example, for the influence model prior, we calculate E c I ( u t , s t | t, c ) =

Original: E c,s t KL ( p ( u t
Revised: E c,s t KL ( p ( u t

Original: | s t , c ) ∥ p ( u t  ))
Revised: | s t , c ) ∥ p ( u t | c ))

Original: ≈ E c,s t KL ( p ( u t
Revised: ⩽ E c,s t KL ( p ( u t

Original: | s t , c ) ∥ (cid:80) ˆ t p ( u t | s′ ˆ t , c )) 
Revised: | s t , c ) ∥ (cid:80) ˆ t p ( u t | s ˆ t , c )) .

Original: .To calculate the input-output mutual information, we use a variant of InfoNCE (van den Oord et al., 2019) estimator.
Revised: This is a variant of InfoNCE (van den Oord et al., 2019) estimator.

Original: That is, we compare the log-probability of the current sample u t ∼ p ( u t | s t , c ) and the log of the average probability of this sample where average is taken over all timesteps within the contiguous batch.
Revised: That is, we compare the log-probability of the current sample u t ∼ p ( u t | s t , c ) and the log of the average probability of this sample where average is taken over all timesteps within the contiguous batch.

Original: This forms a lower bound on true mutual information.
Revised: This forms an upper bound on true mutual information.

Original: In Figure 3 we show the comparison of  I ( u t ,( s t, c t ) | t ) ,  I ( s ′ t , ( s ′ t − 1 , u t ) | t ) ,  I ( s t , ( s t − 1 , a t − 1 ) | t ) , i.e., we do not average over timesteps.
Revised: In Figure 3 we show the comparison of E c I ( u t , s t | t, c ) , E c I ( s ′ t , ( s ′ t − 1 , u t ) | t, c ) , E c I ( s t , ( s t − 1 , a t − 1 ) | t, c ) , i.e., we do not average over timesteps.

Original: The metric growth in the beginning of each plot cannot be explained by the burn-in effects of the neural network distribution because we calculate this metric in a bootstrapping fashion.
Revised: The metric growth in the beginning of each plot cannot be explained by the burn-in effects of the neural network distribution because we calculate this metric in a bootstrapping fashion.

Original: This means, for each timestep, the metric value is obtained given a varying number of present previous timesteps (as each batch contains a short contiguous chunk of the episode), so the burn-in effect would result in a much higher variance.
Revised: This means, for each timestep, the metric value is obtained given a varying number of present previous timesteps (as each batch contains a short contiguous chunk of the episode), so the burn-in effect would result in a much higher variance.

Original: The mutual information metrics show that the prior part of the world model flows the activation of prior distributions in a bottom-up fashion corresponding to the direct influence of cause on effect.
Revised: The mutual information metrics show that the prior part of the world model flows the activation of prior distributions in a bottom-up fashion corresponding to the direct influence of cause on effect.

Original: The robot state MI is activated during the period of the active robot movement, as shown in the middle of Figure 3.
Revised: The robot state MI is activated during the period of the active robot movement, as shown in the middle of Figure 3.

Original: This, in turn, pumps up the MI of the influence latent and the object state.
Revised: This, in turn, pumps up the MI of the influence latent and the object state.

Original: The vertical dashed line indicates the event when the robot starts moving the drawer.
Revised: The vertical dashed line indicates the event when the robot starts moving the drawer.

Original: The object latent mutual information is maximized at some moment after the drawer starts to move.
Revised: The object latent mutual information is maximized at some moment after the drawer starts to move.

Original: We hypothesize that the described behavior can be explained by reactive hierarchical dependencies between the robot state, the influence factor, and the object state that reflects a causal relationship between them.
Revised: We hypothesize that the described behavior can be explained by reactive hierarchical dependencies between the robot state, the influence factor, and the object state that reflects a causal relationship between them.

Original: In other words, each individual MI value does not express the nature of causal dependency, but a combination of them, their co-evolution in time plays as an indicator of correct causal dependencies.
Revised: In other words, each individual MI value does not express the nature of causal dependency, but a combination of them, their co-evolution in time plays as an indicator of correct causal dependencies.

Original: This is also supported by the fact that the prior distributions reflect this property so they produce the excitement given only sequences of the robot actions as a common cause.
Revised: This is also supported by the fact that the prior distributions reflect this property so they produce the excitement given only sequences of the robot actions as a common cause.

Original: In our experiments, we work with multitask environments parameterized by task context c .
Revised: In our experiments, we work with multitask environments parameterized by task context c .

Original: This is done for two reasons: to additionally ensure that the metric behavior is reproducible between the tasks and to ground the work on generalization in task learning in RL, which we aim at for the future work on this model.
Revised: This is done for two reasons: to additionally ensure that the metric behavior is reproducible between the tasks and to ground the work on generalization in task learning in RL, which we aim at for the future work on this model.

Original:  D ISCUSSION AND C ONCLUSION
Revised:  D ISCUSSION AND C ONCLUSION

Original: In this work, we presented CEMA, a structured world model that is capable of multitask reinforcement learning via correct modeling of cause-effect relationships.
Revised: In this work, we presented CEMA, a structured world model that is capable of multitask reinforcement learning via correct modeling of cause-effect relationships.

Original: The model maintains the state of the actuator, the influence of the actuator on the object, and the object state.
Revised: The model maintains the state of the actuator, the influence of the actuator on the object, and the object state.

Original: Using a mutual information estimator, we showed that in proposed model the information flows from cause to effect state as long as the interaction happens.
Revised: Using a mutual information estimator, we showed that in proposed model the information flows from cause to effect state as long as the interaction happens.

Original: The main limitation of the model is that the cause and effect should be manually encoded into the architecture of this model, making it unable to discover causality.
Revised: The main limitation of the model is that the cause and effect should be manually encoded into the architecture of this model, making it unable to discover causality.

Original: Also, we do not assume any internal structure of the space of objects, i.e., the whole set of objects is modeled via one latent.
Revised: Also, we do not assume any internal structure of the space of objects, i.e., the whole set of objects is modeled via one latent.

Original: The proposed model shows the principal way how the generalization with a world model can be achieved: through structural inductive bias, the compression of the information, and restricting models to act on limited input.
Revised: The proposed model shows the principal way how the generalization with a world model can be achieved: through structural inductive bias, the compression of the information, and restricting models to act on limited input.

Original: We plan to research these ideas in the future work on this model.
Revised: We plan to research these ideas in the future work on this model.

Original: Another prominent direction of improvement is eliminating the ground truth segmentation assumption.
Revised: Another prominent direction of improvement is eliminating the ground truth segmentation assumption.

Original: It could additionally reveal the inferred causal structure, for example, if an attention would be used so the model could learn which parts of image to attend to.
Revised: It could additionally reveal the inferred causal structure, for example, if an attention would be used so the model could learn which parts of image to attend to.

Original: This can be an additional sign of correct causal modeling.
Revised: This can be an additional sign of correct causal modeling.

Original: A D ETAILS OF IMPLEMENTATION
Revised: A D ETAILS OF IMPLEMENTATION

Original: Rotated Drawer Environment For all tasks, the drawer’s handle is oriented to the coordinate center, which coincides with the initial position of the drawer’s gripper.
Revised: Rotated Drawer Environment For all tasks, the drawer’s handle is oriented to the coordinate center, which coincides with the initial position of the drawer’s gripper.

Original: The virtual camera is placed straight above the coordinate center so all drawer positions can be viewed from the same angle in the observation.
Revised: The virtual camera is placed straight above the coordinate center so all drawer positions can be viewed from the same angle in the observation.

Original: The action space is the same as in MetaWorld, a four-dimensional vector representing shift vector for the gripper and normalized torque applied to the gripper fingers.
Revised: The action space is the same as in MetaWorld, a four-dimensional vector representing shift vector for the gripper and normalized torque applied to the gripper fingers.

Original: The observation is a 64 × 64 RGB image.
Revised: The observation is a 64 × 64 RGB image.

Original: Examples of the image observations are in Figure 2.
Revised: Examples of the image observations are in Figure 2.

Original: The reward function is defined in the same way as in the DrawerOpen-v2 task of MetaWorld (Yu et al., 2019).
Revised: The reward function is defined in the same way as in the DrawerOpen-v2 task of MetaWorld (Yu et al., 2019).

Original: Namely, the reward is normalized in the range [0 , 10] where maximal reward corresponds to the solved task.
Revised: Namely, the reward is normalized in the range [0 , 10] where maximal reward corresponds to the solved task.

Original: The length of each episode is equal to 500 .
Revised: The length of each episode is equal to 500 .

Original: Agent architecture.
Revised: Agent architecture.

Original: We use the ground truth actor-object segmentation provided by Metaworld environment.
Revised: We use the ground truth actor-object segmentation provided by Metaworld environment.

Original: Also, to provide the model with an unoccluded image of the object, we render an additional image with a transparent robot.
Revised: Also, to provide the model with an unoccluded image of the object, we render an additional image with a transparent robot.

Original: This way, we naturally avoid an undesired behavior of the model via robot occlusion understanding.
Revised: This way, we naturally avoid an undesired behavior of the model via robot occlusion understanding.

Original: Actor RSSM implementation with all its parameters is identical to Dreamer.
Revised: Actor RSSM implementation with all its parameters is identical to Dreamer.

Original: All distributions in the model are assumed to be Gaussian.
Revised: All distributions in the model are assumed to be Gaussian.

Original: Both prior and posterior distributions of the influence factor are represented by a single-layered MLP with 200 units.
Revised: Both prior and posterior distributions of the influence factor are represented by a single-layered MLP with 200 units.

Original: The object layer has the same imagination mechanism as in Dreamer, but instead of actions, it is conditioned on the influence factor.
Revised: The object layer has the same imagination mechanism as in Dreamer, but instead of actions, it is conditioned on the influence factor.

Original: Also, we use the deterministic part of the object state only in imagination and disregard it anywhere else since we assume that most of the information is contained in the current image.
Revised: Also, we use the deterministic part of the object state only in imagination and disregard it anywhere else since we assume that most of the information is contained in the current image.

Original: Instead of predicting the full image, separate decoders for object and actor images are used.
Revised: Instead of predicting the full image, separate decoders for object and actor images are used.

Original: Changes were made to the training process of the object layer to account for the specifics of the object layer.
Revised: Changes were made to the training process of the object layer to account for the specifics of the object layer.

Original: In particular, the free KL divergence gap between object prior and posterior originally introduced in the Dreamer algorithm, was removed, and the corresponding loss weight was increased by three times.
Revised: In particular, the free KL divergence gap between object prior and posterior originally introduced in the Dreamer algorithm, was removed, and the corresponding loss weight was increased by three times.

Original: We justify these changes with three following reasons: to account for low stochasticity of object appearance, to employ ideas from
Revised: We justify these changes with three following reasons: to account for low stochasticity of object appearance, to employ ideas from

Original: RPC (Eysenbach et al., 2021), and to additionally push prior distribution towards posterior, since learning influence vector through completely stochastic object state is difficult for the model.
Revised: RPC (Eysenbach et al., 2021), and to additionally push prior distribution towards posterior, since learning influence vector through completely stochastic object state is difficult for the model.

Original: B RL A GENTS P ERFORMANCE
Revised: B RL A GENTS P ERFORMANCE

Revised: C M UTUAL I NFORMATION U PPER B OUND D ERIVATION

Revised: For example, consider an influence factor prior distribution p ( u t | s t , c ) :

Revised: [Equation]

Revised: This way we recovered an upper bound for the task-averaged mutual information.

Revised: For other distributions, all latent factors should be explicitly conditioned on the task context c .

Revised: Effectively, this task context may be omitted in the final formulas given the structure of our graphical model.

Revised: This is an upper bound since we approximate the marginal distribution in the RHS of KL.

Revised: Our estimator simply averages probability scores of a u t sample over all distributions in time.

Revised: The bound has the form similar to InfoNCE (van