Original: L EARNING O BJECT -O RIENTED D YNAMICS FOR P LANNING FROM T EXT
Revised: L EARNING O BJECT -O RIENTED D YNAMICS FOR P LANNING FROM T EXT

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: The advancement of dynamics models enables model-based planning in complex environments.
Revised: The advancement of dynamics models enables model-based planning in complex environments.

Original: Dynamics models mostly study image-based games with fully observable states.
Revised: Dynamics models mostly study image-based games with fully observable states.

Original: Generalizing these models to Text-Based Games (TBGs), which often include partially observable states with noisy text observations, is challenging.
Revised: Generalizing these models to Text-Based Games (TBGs), which often include partially observable states with noisy text observations, is challenging.

Original: In this work, we propose an Object-Oriented Text Dynamics (OOTD) model that enables planning algorithms to solve decision-making problems in text domains.
Revised: In this work, we propose an Object-Oriented Text Dynamics (OOTD) model that enables planning algorithms to solve decision-making problems in text domains.

Original: OOTD predicts a memory graph that dynamically remembers the history of object observations and ﬁlters object-irrelevant information.
Revised: OOTD predicts a memory graph that dynamically remembers the history of object observations and ﬁlters object-irrelevant information.

Original: To improve the robustness of dynamics, our OOTD model identiﬁes the objects inﬂuenced by input actions and predicts beliefs of object states with independently parameterized transition layers.
Revised: To improve the robustness of dynamics, our OOTD model identiﬁes the objects inﬂuenced by input actions and predicts beliefs of object states with independently parameterized transition layers.

Original: We develop variational objectives under the object-supervised and self-supervised settings to model the stochasticity of predicted dynamics.
Revised: We develop variational objectives under the object-supervised and self-supervised settings to model the stochasticity of predicted dynamics.

Original: Empirical results show that our OOTD-based planner signiﬁcantly outperforms model-free baselines in terms of sample efﬁciency and running scores.
Revised: Empirical results show that our OOTD-based planner signiﬁcantly outperforms model-free baselines in terms of sample efﬁciency and running scores.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Planning algorithms typically leverage the environment dynamics to solve decision-making problems (Sutton & Barto, 2018).
Revised: Planning algorithms typically leverage the environment dynamics to solve decision-making problems (Sutton & Barto, 2018).

Original: To plan in unknown environments, the agent must learn a dynamics model to predict future (belief) states and rewards by conditioning on an action.
Revised: To plan in unknown environments, the agent must learn a dynamics model to predict future (belief) states and rewards by conditioning on an action.

Original: This dynamics model enables the implementation of intensive search for optimal actions, which can potentially increase both sample efﬁciency and cumulative rewards, compared to model-free methods (Hafner et al., 2019;
Revised: This dynamics model enables the implementation of intensive search for optimal actions, which can potentially increase both sample efﬁciency and cumulative rewards, compared to model-free methods (Hafner et al., 2019;

Original: Wang et al., 2019; Kaiser et al., 2020).
Revised: Wang et al., 2019; Kaiser et al., 2020).

Original: Despite the promising performance, learning a dynamics model that can accurately generalize at test time is still challenging, especially when handling a high-dimensional state space for low-level features, e.g., pixels and text.
Revised: Despite the promising performance, learning a dynamics model that can accurately generalize at test time is still challenging, especially when handling a high-dimensional state space for low-level features, e.g., pixels and text.

Original: To facilitate dynamics learning in complex environments, Diuk et al.
Revised: To facilitate dynamics learning in complex environments, Diuk et al.

Original: (2008) proposed an ObjectOriented Markov Decision Process (OO-MDP) that factorizes world states into object states.
Revised: (2008) proposed an ObjectOriented Markov Decision Process (OO-MDP) that factorizes world states into object states.

Original: They showed that the agent can ﬁnd optimal policies with a better sample efﬁciency by modeling the dynamics at the object level.
Revised: They showed that the agent can ﬁnd optimal policies with a better sample efﬁciency by modeling the dynamics at the object level.

Original: Some following works (Finn et al., 2016; Goel et al., 2018; Zhu et al., 2018; 2020) extended OO-MDPs to image-based games.
Revised: Some following works (Finn et al., 2016; Goel et al., 2018; Zhu et al., 2018; 2020) extended OO-MDPs to image-based games.

Original: These methods typically assume full observability over game states and a ﬁxed input size, which facilitates the use of object masks to decompose an image into different objects.
Revised: These methods typically assume full observability over game states and a ﬁxed input size, which facilitates the use of object masks to decompose an image into different objects.

Original: On the other hand, in Text-Based Games (TBGs), the text observation, whose length runs from a few words to an entire paragraph, is a partial description of the current game state, and thus each observation provides information about a limited number of objects (in fact, only an average of 4 .
Revised: On the other hand, in Text-Based Games (TBGs), the text observation, whose length runs from a few words to an entire paragraph, is a partial description of the current game state, and thus each observation provides information about a limited number of objects (in fact, only an average of 4 .

Original: 51% of the candidate objects are mentioned in an observation from Textworld (Côté et al., 2018)).
Revised: 51% of the candidate objects are mentioned in an observation from Textworld (Côté et al., 2018)).

Original: The dynamics model must remember the history of observed objects to predict accurately their states.
Revised: The dynamics model must remember the history of observed objects to predict accurately their states.

Original: Moreover, an observation typically contains lots of noisy patterns that record object-irrelevant information (e.g., the non-bolded text from o t in Figure 1).
Revised: Moreover, an observation typically contains lots of noisy patterns that record object-irrelevant information (e.g., the non-bolded text from o t in Figure 1).

Original: While previous works (Ammanabrolu & Riedl, 2019; Ammanabrolu & Hausknecht, 2020) designed rule-based heuristics to extract useful information, we expect that dynamically capturing object information from noisy observations helps to generalize to more environments.
Revised: While previous works (Ammanabrolu & Riedl, 2019; Ammanabrolu & Hausknecht, 2020) designed rule-based heuristics to extract useful information, we expect that dynamically capturing object information from noisy observations helps to generalize to more environments.

Original: In this work, we design an Object-Oriented Text Dynamics (OOTD ) model that integrates: 1) Graph Representation for Objects.
Revised: In this work, we design an Object-Oriented Text Dynamics (OOTD ) model that integrates: 1) Graph Representation for Objects.

Original: OOTD predicts a memory graph from the input states.
Revised: OOTD predicts a memory graph from the input states.

Original: The memory graph, whose nodes correspond to the objects to be modelled, captures the information about these objects from the beginning of a game.
Revised: The memory graph, whose nodes correspond to the objects to be modelled, captures the information about these objects from the beginning of a game.

Original: By applying this graph as an information bottleneck for dynamics prediction, OOTD ﬁlters the object-irrelevant information in the text inputs, which forms a tighter information plane and facilitates generalization to unseen games (Tishby & Zaslavsky, 2015).
Revised: By applying this graph as an information bottleneck for dynamics prediction, OOTD ﬁlters the object-irrelevant information in the text inputs, which forms a tighter information plane and facilitates generalization to unseen games (Tishby & Zaslavsky, 2015).

Original: 2) Independent Transition Layers.
Revised: 2) Independent Transition Layers.

Original: Intuitively, an input action inﬂuences the states of a limited number of objects, for example, when the agent performs " take the carrot from the counter ", only the states of " carrot " and " counter " should be updated.
Revised: Intuitively, an input action inﬂuences the states of a limited number of objects, for example, when the agent performs " take the carrot from the counter ", only the states of " carrot " and " counter " should be updated.

Original: To capture which objects are inﬂuenced and keep others invariant to this update, OOTD learns an action-aware representation for each object with bi-directional attention and predicts the belief states of objects with independently parameterized transition layers.
Revised: To capture which objects are inﬂuenced and keep others invariant to this update, OOTD learns an action-aware representation for each object with bi-directional attention and predicts the belief states of objects with independently parameterized transition layers.

Original: This independent mechanism enhances the robustness for dynamics prediction (Goyal et al., 2019)
Revised: This independent mechanism enhances the robustness for dynamics prediction (Goyal et al., 2019)

Original: (also see the ablation study in Section 4.2).
Revised: (also see the ablation study in Section 4.2).

Original: During planning, we predict belief states without knowing rewards and observations.
Revised: During planning, we predict belief states without knowing rewards and observations.

Original: To learn a stochastic dynamics model, we propose object-supervised and self-supervised Evidence Lower Bound (ELBo) objectives for training our OOTD model.
Revised: To learn a stochastic dynamics model, we propose object-supervised and self-supervised Evidence Lower Bound (ELBo) objectives for training our OOTD model.

Original: We evaluate how well the OOTD model supports planning by implementing planning algorithms (e.g., Dyna-Q (Kuvayev & Sutton, 1996) and Monte
Revised: We evaluate how well the OOTD model supports planning by implementing planning algorithms (e.g., Dyna-Q (Kuvayev & Sutton, 1996) and Monte

Original: Carlo Tree Search (MCTS) (Kocsis & Szepesvári, 2006))
Revised: Carlo Tree Search (MCTS) (Kocsis & Szepesvári, 2006))

Original: based on learned dynamics.
Revised: based on learned dynamics.

Original: Empirical results show that these planning algorithms signiﬁcantly outperform other baselines in terms of game performance and sample efﬁciency.
Revised: Empirical results show that these planning algorithms signiﬁcantly outperform other baselines in terms of game performance and sample efﬁciency.

Original: To support the design of our OOTD model, we include an ablation study that visualizes the object states and quantiﬁes the dynamics prediction performance.
Revised: To support the design of our OOTD model, we include an ablation study that visualizes the object states and quantiﬁes the dynamics prediction performance.

Original: Contributions: 1) We introduce our approach to implementing model-based planning algorithms to solve decision-making problems in the text-domain.
Revised: Contributions: 1) We introduce our approach to implementing model-based planning algorithms to solve decision-making problems in the text-domain.

Original: We propose the OOTD model that learns to predict beliefs about object dynamics.
Revised: We propose the OOTD model that learns to predict beliefs about object dynamics.

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: meal . ingredients : banana , block of cheese ,
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original:  directions : dice the banana,fry the banana,chop the block of cheese , roast the block of cheese , slice the carrot,fry the [Figure] , and preparemeal.
Revised: ingredients : banana , block of cheese , carrot directions : dice the banana,fry the banana,chop the block of cheese , roast the block of cheese , slice the carrot,fry the carrot , and preparemeal.

Original: Take banana from counter 𝑒𝑛𝑣
Revised: Take banana from counter 𝑒𝑛𝑣

Original: Encode 𝑜 " and 𝑎 "’( With { 𝑞 /( , 𝑞 /+ , 𝑞 /0 }
Revised: Encode 𝑜 " and 𝑎 "’( With { 𝑞 /( , 𝑞 /+ , 𝑞 /0 }

Original: Plan with Transition( 𝑝 𝒯( , 𝑝 𝒯+ , 𝑝 𝒯0 ) andreward( 𝑝 3 ) Models
Revised: Plan with Transition( 𝑝 𝒯( , 𝑝 𝒯+ , 𝑝 𝒯0 ) andreward( 𝑝 3 ) Models

Original: Controllingin TBG
Revised: Controllingin TBG

Original: Planninginobject space
Revised: Planninginobject space

Original: Examine cookbook 𝒛 (,"’( 𝒛 0,"’(
Revised: Examine cookbook 𝒛 (,"’( 𝒛 0,"’(

Original: Planner 𝒛 +," 𝒛 (," 𝒛 0," 𝒛 +,")
Revised: Planner 𝒛 +," 𝒛 (," 𝒛 0," 𝒛 +,")

Original: ( 𝒛 (,")( 𝒛 0,")(
Revised: ( 𝒛 (,")( 𝒛 0,")(

Original: You take the banana from the counter .
Revised: You take the banana from the counter .

Original: [Figure]
Revised: [Figure]

Original:  O BJECT L EVEL P LANNING IN T EXT D OMAINS
Revised:  O BJECT L EVEL P LANNING IN T EXT D OMAINS

Original: In this work, the object-level dynamics are learned and evaluated in Text-Based Games (TBGs), for which we extend OO-MDPs (Diuk et al., 2008) to capture partial observability.
Revised: In this work, the object-level dynamics are learned and evaluated in Text-Based Games (TBGs), for which we extend OO-MDPs (Diuk et al., 2008) to capture partial observability.

Original: We deﬁne a graph representation to remember object information and introduce the corresponding planning algorithms.
Revised: We deﬁne a graph representation to remember object information and introduce the corresponding planning algorithms.

Original: 2. T EXT -B ASED G AMES
Revised: 2. T EXT -B ASED G AMES

Original: Text-Based Games (TBGs) are complex, interactive simulations where an agent explores and plays a game by processing text observations and responding with text commands.
Revised: Text-Based Games (TBGs) are complex, interactive simulations where an agent explores and plays a game by processing text observations and responding with text commands.

Original: To study the generalization ability of this policy, many recent works (Yin & May, 2019; Hausknecht et al., 2019; Ecoffet et al., 2019; Adolphs & Hofmann, 2020; Adhikari et al., 2020) proposed to learn a policy by training with a set of games and evaluating with games from a hold-out test set.
Revised: To study the generalization ability of this policy, many recent works (Yin & May, 2019; Hausknecht et al., 2019; Ecoffet et al., 2019; Adolphs & Hofmann, 2020; Adhikari et al., 2020) proposed to learn a policy by training with a set of games and evaluating with games from a hold-out test set.

Original: The distributions of rewards over state-action pairs are different in training and testing games because TBGs assign rewards by conditioning on speciﬁc goals (that are predeﬁned, but unknown).
Revised: The distributions of rewards over state-action pairs are different in training and testing games because TBGs assign rewards by conditioning on speciﬁc goals (that are predeﬁned, but unknown).

Original: In this setting, TBGs have different intermediate goals in each game, but they commonly share the same ultimate goals, and the underlying causal dependencies leading to this ultimate goal are consistent.
Revised: In this setting, TBGs have different intermediate goals in each game, but they commonly share the same ultimate goals, and the underlying causal dependencies leading to this ultimate goal are consistent.

Original: For example, to solve the ‘First Text-World Problems’ (Côté et al., 2018), an agent should always gather and process cooking ingredients (e.g., " fry the potato ") according to a recipe that it discovers in the game, although the names and locations of the ingredients are different across different games.
Revised: For example, to solve the ‘First Text-World Problems’ (Côté et al., 2018), an agent should always gather and process cooking ingredients (e.g., " fry the potato ") according to a recipe that it discovers in the game, although the names and locations of the ingredients are different across different games.

Original: This property enables the generalization of knowledge learned from training games to solve testing games.
Revised: This property enables the generalization of knowledge learned from training games to solve testing games.

Original: 2. O BJECT -O RIENTED P ARTIALLY O BSERVABLE M ARKOV D ECISION P ROCESSES
Revised: 2. O BJECT -O RIENTED P ARTIALLY O BSERVABLE M ARKOV D ECISION P ROCESSES

Original: Since an observation reveals only partial information about the game state, we formulate TBGs as partially observable Markov decision processes (POMDPs).
Revised: Since an observation reveals only partial information about the game state, we formulate TBGs as partially observable Markov decision processes (POMDPs).

Original: In this work, we extend the POMDP framework to represent object-level dynamics (Diuk et al., 2008).
Revised: In this work, we extend the POMDP framework to represent object-level dynamics (Diuk et al., 2008).

Original: The Object-Oriented (OO) POMDP is a tuple (cid:104)S , O , Z , Φ , G , A , R , T , γ (cid:105) , where:
Revised: The Object-Oriented (OO) POMDP is a tuple (cid:104)S , O , Z , Φ , G , A , R , T , γ (cid:105) , where:

Original: • S and O deﬁne the spaces of low-level states and observations from the TGB environments while Z and Φ are the spaces of object-level states and observations.
Revised: • S and O deﬁne the spaces of low-level states and observations from the TGB environments while Z and Φ are the spaces of object-level states and observations.

Original: In a TBG, o ∈ O is a text message from the game environment (e.g., " Please fry the potato. "), whereas φ ∈ Φ records the speciﬁc objects and their relations in this message (e.g., a triplet like " potato-need-fry ")).
Revised: In a TBG, o ∈ O is a text message from the game environment (e.g., " Please fry the potato. "), whereas φ ∈ Φ records the speciﬁc objects and their relations in this message (e.g., a triplet like " potato-need-fry ")).

Original: To model the object dynamics, the agent must extract φ from o by distilling objects’ information from text sentences.
Revised: To model the object dynamics, the agent must extract φ from o by distilling objects’ information from text sentences.

Original: In this work, we model object states z ∈ Z (instead of s ∈ S ) and learn latent representations for z .
Revised: In this work, we model object states z ∈ Z (instead of s ∈ S ) and learn latent representations for z .

Original: • G and A are the spaces for goals and actions.
Revised: • G and A are the spaces for goals and actions.

Original: We study choice-based games (Yin & May, 2019), where candidate actions A t ∈ A are available at time step t .
Revised: We study choice-based games (Yin & May, 2019), where candidate actions A t ∈ A are available at time step t .

Original: We include a goal variable g to mark different tasks in each game.
Revised: We include a goal variable g to mark different tasks in each game.

Original: Following the Universal MDP (UMDP) (Schaul et al., 2015), the agent initializes a goal at the beginning of a game and updates it when the task is ﬁnished.
Revised: Following the Universal MDP (UMDP) (Schaul et al., 2015), the agent initializes a goal at the beginning of a game and updates it when the task is ﬁnished.

Original: • R and T deﬁne the spaces of reward and transition models.
Revised: • R and T deﬁne the spaces of reward and transition models.

Original: γ is a discount factor.
Revised: γ is a discount factor.

Original: We assume the real dynamics models ( p ∗T and p ∗ r ) are unknown, so we learn the object-oriented transition model and the reward model, i.e., ∀ k ∈ { 1 , . .
Revised: We assume the real dynamics models ( p ∗T and p ∗ r ) are unknown, so we learn the object-oriented transition model and the reward model, i.e., ∀ k ∈ { 1 , . .

Original: , K } , p k T ( z k,t | z t − 1 , a t − 1 ) ∈ T and p r ( r t | 
Revised: , K } , p k T ( z k,t | z t − 1 , a t − 1 ) ∈ T and p r ( r t | z

Original: z t − 1 , g t ) ∈ R .
Revised: t − 1 , g t ) ∈ R .

Original: In our OO-POMDP, transitions p T , observations φ , and states z are modelled for each object, whereas actions a and rewards r are deﬁned for the entire environment.
Revised: In our OO-POMDP, transitions p T , observations φ , and states z are modelled for each object, whereas actions a and rewards r are deﬁned for the entire environment.

Original: This generalizes to popular RL environments that accept an action and return a reward at every time step.
Revised: This generalizes to popular RL environments that accept an action and return a reward at every time step.

Original: Unlike (Wandzel et al., 2019) that represents object states with symbolic attributes, we develop a latent object representation that can generalize to the complex environment with high-dimensional inputs (e.g., text).
Revised: Unlike (Wandzel et al., 2019) that represents object states with symbolic attributes, we develop a latent object representation that can generalize to the complex environment with high-dimensional inputs (e.g., text).

Original: 2. G RAPH R EPRESENTATION FOR O BJECTS
Revised: G RAPH R EPRESENTATION FOR O BJECTS

Original: We utilize an object-relation-object triplet φ to represent object information.
Revised: We utilize an object-relation-object triplet φ to represent object information.

Original: At time step t, φ t =
Revised: At time step t, φ t =

Original: [ φ t, 1 , ..., φ t,M ] ( M is the number of observed triplets at t ).
Revised: [ φ t, 1 , ..., φ t,M ] ( M is the number of observed triplets at t ).

Original: Given a total of K candidate objects and
Revised: Given a total of K candidate objects and

Original: C candidate relations, these triplets can be mapped to a knowledge graph Ω t ∈ { 0 , 1 } C × K × K .
Revised: C candidate relations, these triplets can be mapped to a knowledge graph Ω t ∈ { 0 , 1 } C × K × K .

Original: Each entry ( c, i, j ) ∈ { 0 , 1 } indicates whether there is a relationship r between the i th and j th objects.
Revised: Each entry ( c, i, j ) ∈ { 0 , 1 } indicates whether there is a relationship r between the i th and j th objects.

Original: This knowledge graph forms a natural representation of object relations since object information in most observations o t corresponds to either entity attributes (e.g., " potato-is-fried ") or to relational information about entities(e.g.
Revised: This knowledge graph forms a natural representation of object relations since object information in most observations o t corresponds to either entity attributes (e.g., " potato-is-fried ") or to relational information about entities(e.g.

Original: , " potato-on-counter " and " bedroom-north_of-kitchen "), which can be conveniently represented by triplets and the corresponding graph.
Revised: , " potato-on-counter " and " bedroom-north_of-kitchen "), which can be conveniently represented by triplets and the corresponding graph.

Original: Memory Graph.
Revised: Memory Graph.

Original: We store the observations up to time step t in a memory graph h t
Revised: We store the observations up to time step t in a memory graph h t

Original: =
Revised: =

Original: Ω 0 ⊕ Ω 1 ⊕ · · · ⊕ Ω t that captures object information observed by the agent since the beginning of a game.
Revised: Ω 0 ⊕ Ω 1 ⊕ · · · ⊕ Ω t that captures object information observed by the agent since the beginning of a game.

Original: Similar to previous works Adhikari et al.
Revised: Similar to previous works Adhikari et al.

Original: (2020); Ammanabrolu & Riedl (2019); Ammanabrolu &
Revised: (2020); Ammanabrolu & Riedl (2019); Ammanabrolu &

Original: Hausknecht (2020), we summarize the object-level history with a latent memory graph h t .
Revised: Hausknecht (2020), we summarize the object-level history with a latent memory graph h t .

Original: To update the memory from t − 1 to t , we learn a graph updater ⊕ , and h t = h t − 1 ⊕ Ω t .
Revised: To update the memory from t − 1 to t , we learn a graph updater ⊕ , and h t = h t − 1 ⊕ Ω t .

Original: During updating, ⊕ needs to resolves some semantically contradictory triplets, for example " player-at-kitchen " and " player-at-bedroom " (because the player cannot simultaneously be in two different locations).
Revised: During updating, ⊕ needs to resolves some semantically contradictory triplets, for example " player-at-kitchen " and " player-at-bedroom " (because the player cannot simultaneously be in two different locations).

Original: Our
Revised: Our

Original: OOTD model is trained to automatically emulate such an operator.
Revised: OOTD model is trained to automatically emulate such an operator.

Original: Our transition model (Section 3.1) is trained to generate h t from z t , allowing the latent object states to capture object relations.
Revised: Our transition model (Section 3.1) is trained to generate h t from z t , allowing the latent object states to capture object relations.

Original: 2. M ODEL -B ASED P LANNING IN T EXT -B ASED G AMES
Revised: 2. M ODEL -B ASED P LANNING IN T EXT -B ASED G AMES

Original: [Table]
Revised: [Table]

Original: We introduce model-based planning in
Revised: We introduce model-based planning in

Original: TBGs.
Revised: TBGs.

Original: Based on the OO-POMDP (Section 2.2), at each time step t , we deﬁne latent states z t =
Revised: Based on the OO-POMDP (Section 2.2), at each time step t , we deﬁne latent states z t =

Original: [ z 1 ,t , . .
Revised: [ z 1 ,t , . .

Original: , z K,t ] of K objects action commands a t , and scalar rewards r t , that follow the stochastic dynamics in
Revised: , z K,t ] of K objects action commands a t , and scalar rewards r t , that follow the stochastic dynamics in

Original: Table  details of these dynamics models are introduced in Section 3.
Revised: Table 1. The details of these dynamics models are introduced in Section 3.

Original: Based on these dynamics, we implement planning algorithms to select an action from candidate commands A t for maximizing the expected sum of rewards E π
Revised: Based on these dynamics, we implement planning algorithms to select an action from candidate commands A t for maximizing the expected sum of rewards E π

Original: ( (cid:80) Tt =

Revised: ( (cid:80) Tt =

Original: γ t r t ) , as shown in Figure 1.
Revised: γ t r t ) , as shown in Figure 1.

Original: In this work, we study choice-based games (Yin & May, 2019), where the candidate commands (or actions)
Revised: In this work, we study choice-based games (Yin & May, 2019), where the candidate commands (or actions)

Original: A t are available and the planner determines the action a t ∈
Revised: A t are available and the planner determines the action a t ∈

Original: A t to be performed.
Revised: A t to be performed.

Original: The planning algorithms include Dyna-Q, Monte-Carlo Tree Search (MCTS) and their combinations:
Revised: The planning algorithms include Dyna-Q, Monte-Carlo Tree Search (MCTS) and their combinations:

Original: Dyna-Q (Kuvayev & Sutton, 1996) incorporates dynamics models and Q-learning.
Revised: Dyna-Q (Kuvayev & Sutton, 1996) incorporates dynamics models and Q-learning.

Original: Dyna-Q interactively 1) updates the dynamics model with observed transitions and 2) trains the Q network to minimize the Temporal Difference (TD) loss (Equation 1) based on both observed transitions from the environment and predicted transitions from the dynamics models.
Revised: Dyna-Q interactively 1) updates the dynamics model with observed transitions and 2) trains the Q network to minimize the Temporal Difference (TD) loss (Equation 1) based on both observed transitions from the environment and predicted transitions from the dynamics models.

Original: Compared to the model-free Deep
Revised: Compared to the model-free Deep

Original: Q-Network (DQN), Dyna-Q is more sample efﬁcient: by expanding the replay buffer with dynamics models, Dyna-Q converges faster with the same number of interactions from the environment.
Revised: Q-Network (DQN), Dyna-Q is more sample efﬁcient: by expanding the replay buffer with dynamics models, Dyna-Q converges faster with the same number of interactions from the environment.

Original: [Equation]
Revised: [Equation]

Original: MCTS (Kocsis & Szepesvári, 2006) is a heuristic search algorithm that builds and updates a search tree based on environment dynamics.
Revised: MCTS (Kocsis & Szepesvári, 2006) is a heuristic search algorithm that builds and updates a search tree based on environment dynamics.

Original: By performing Monte-Carlo simulations organized in a tree structure, MCTS does an efﬁcient search in environments with large action spaces (Silver et al., 2018).
Revised: By performing Monte-Carlo simulations organized in a tree structure, MCTS does an efﬁcient search in environments with large action spaces (Silver et al., 2018).

Original: MCTS iteratively runs multiple playouts.
Revised: MCTS iteratively runs multiple playouts.

Original: At the i th playout, we implement 1) Selection: Traverse the tree from the root node to a leaf node (corresponding to object state z τ ) by selecting the action command a i,t to maximize the Upper Conﬁdence Bound (UCB) (Couëtoux et al., 2011):
Revised: At the i th playout, we implement 1) Selection: Traverse the tree from the root node to a leaf node (corresponding to object state z τ ) by selecting the action command a i,t to maximize the Upper Conﬁdence Bound (UCB) (Couëtoux et al., 2011):

Original: [Equation]
Revised: [Equation]

Original: where c puct controls the scale of exploration, and ζ i records the visit count at the i th play.
Revised: where c puct controls the scale of exploration, and ζ i records the visit count at the i th play.

Original: Evaluation: Evaluate the selected leaf
Revised: Evaluation: Evaluate the selected leaf

Original: z τ with the reward model p r
Revised: z τ with the reward model p r

Original: ( r τ | z τ , g t ) .
Revised: ( r τ | z τ , g t ) .

Original: 3) Expansion: Expand the leaf node by adding child nodes.
Revised: 3) Expansion: Expand the leaf node by adding child nodes.

Original: Back Up: Update the action-values: Q i =
Revised: Back Up: Update the action-values: Q i =

Original: ( Q i − 1 ζ i − 1 + r τ ) /
Revised: ( Q i − 1 ζ i − 1 + r τ ) /

Original: ( ζ i − 1 + 1) and increment the visit count: ζ i = ζ
Revised: ( ζ i − 1 + 1) and increment the visit count: ζ i = ζ

Original: i − 1 + 1 on all the traversed edges.
Revised: i − 1 + 1 on all the traversed edges.

Original: Dyna-Q + MCTS initializes Q values in MCTS by utilizing the Dyna-Q network.
Revised: Dyna-Q + MCTS initializes Q values in MCTS by utilizing the Dyna-Q network.

Original: This combination enables MCTS to perform the tree search with more efﬁciency and efﬁcacy (Silver et al., 2016).
Revised: This combination enables MCTS to perform the tree search with more efﬁciency and efﬁcacy (Silver et al., 2016).

Original:  O BJECT -O RIENTED D YNAMICS M ODELS
Revised:  O BJECT -O RIENTED D YNAMICS M ODELS

Original: We introduce the Object-Oriented Text Dynamics (OOTD) models, including 1) a transition model that directly predicts the belief states of objects and 2) a reward model that maps sampled object states to rewards.
Revised: We introduce the Object-Oriented Text Dynamics (OOTD) models, including 1) a transition model that directly predicts the belief states of objects and 2) a reward model that maps sampled object states to rewards.

Original: Based on these models, the planner determines an action purely based on predicted beliefs and calibrates its belief with the feedback (observations and rewards) from the environment.
Revised: Based on these models, the planner determines an action purely based on predicted beliefs and calibrates its belief with the feedback (observations and rewards) from the environment.

Original: We introduce the object-supervised and self-supervised objectives to train these models.
Revised: We introduce the object-supervised and self-supervised objectives to train these models.

Original: RGCNGraphEncoder
Revised: [Figure]

Original: BIDAFNetwork 𝒂 "#$ 𝒉 "#$
Revised: BIDAFNetwork 𝒂 "#$ 𝒉 "#$

Original: MLP
Revised: MLP

Original: MLP
Revised: MLP

Original: MLP 𝒆 $ 𝒆 ’ 𝒆 ( …… …… ComplExGraphDecoder 𝒛 $,"#$ 𝒛 ’,"#$ 𝒛 (,"#$ …… ……
Revised: MLP 𝒆 $ 𝒆 ’ 𝒆 ( …… …… ComplExGraphDecoder 𝒛 $,"#$ 𝒛 ’,"#$ 𝒛 (,"#$ …… ……

Original: [Figure]
Revised: [Figure]

Original: 3. O BJECT O RIENTED T RANSITION M ODEL
Revised: 3. O BJECT O RIENTED T RANSITION M ODEL

Original: Figure 2 illustrates our transition model, which has three major components (a) a graph decoder, (b) a graph encoder, and (c) independent transition layers.
Revised: Figure 2 illustrates our transition model, which has three major components (a) a graph decoder, (b) a graph encoder, and (c) independent transition layers.

Original: (a) Graph Decoder.
Revised: (a) Graph Decoder.

Original: Given the states of K objects z t − 1 =
Revised: Given the states of K objects z t − 1 =

Original: [ z 1 ,t − 1 , ..., z K,t − 1 ] , we implement a graph decoder with the ComplEx scoring function (Trouillon et al., 2016) that maps the object states to a memory graph h t − 1 ∈ [0 , 1] C × K × K .
Revised: [ z 1 ,t − 1 , ..., z K,t − 1 ] , we implement a graph decoder with the ComplEx scoring function (Trouillon et al., 2016) that maps the object states to a memory graph h t − 1 ∈ [0 , 1] C × K × K .

Original: The scoring function for an entry ( c, i, j ) (in the adjacency matrix of h t − 1 ) is implemented by Re ( (cid:104) ω
Revised: The scoring function for an entry ( c, i, j ) (in the adjacency matrix of h t − 1 ) is implemented by Re ( (cid:104) ω

Original: c , z i,t − 1 , z j,t − 1 (cid:105) ) where ω c ∈ C E is a complex vector for the relation c and Re denotes the real part of the decomposition in complex space.
Revised: c , z i,t − 1 , z j,t − 1 (cid:105) ) where ω c ∈ C E is a complex vector for the relation c and Re denotes the real part of the decomposition in complex space.

Original: The function can model a large variety of binary relations, including both symmetric and antisymmetric relations.
Revised: The function can model a large variety of binary relations, including both symmetric and antisymmetric relations.

Original: Based on this scoring function, predicting h t − 1 is equivalent to estimating the probability of having a relation c between object i and j , which can be efﬁciently approximated with a low-ranked ( E (cid:28) K ) matrix decomposition: h t =
Revised: Based on this scoring function, predicting h t − 1 is equivalent to estimating the probability of having a relation c between object i and j , which can be efﬁciently approximated with a low-ranked ( E (cid:28) K ) matrix decomposition: h t =

Original: [ Re ( ZW 1 Z T ) , ..., Re ( Z W C Z T )] where Z ∈ R K × E is the matrix of object states and W c ∈ C E × E is a complex matrix.
Revised: [ Re ( ZW 1 Z T ) , ..., Re ( Z W C Z T )] where Z ∈ R K × E is the matrix of object states and W c ∈ C E × E is a complex matrix.

Original: In practice, this decomposition enables to map efﬁciently from z t − 1 to h t − 1 , therefore scaling our method to environments with numerous objects.
Revised: In practice, this decomposition enables to map efﬁciently from z t − 1 to h t − 1 , therefore scaling our method to environments with numerous objects.

Original: (b) Graph Encoder.
Revised: (b) Graph Encoder.

Original: We encode the memory graph h t − 1 into node representations e t − 1 =
Revised: We encode the memory graph h t − 1 into node representations e t − 1 =

Original: [ e 1 ,t − 1 , ..., e K,t − 1 ] with a Relational Graph Convolution Network (R-GCN) (Schlichtkrull et al., 2018).
Revised: [ e 1 ,t − 1 , ..., e K,t − 1 ] with a Relational Graph Convolution Network (R-GCN) (Schlichtkrull et al., 2018).

Original: Nodes in h t correspond to candidate objects and RGCN performs message passing between nodes, so e k,t − 1 captures the semantic information of the k th object and its relations to other nodes (objects).
Revised: Nodes in h t correspond to candidate objects and RGCN performs message passing between nodes, so e k,t − 1 captures the semantic information of the k th object and its relations to other nodes (objects).

Original: Compared to z t − 1 , e t − 1 denotes the object embeddings after taking into account relations to other objects in the graph.
Revised: Compared to z t − 1 , e t − 1 denotes the object embeddings after taking into account relations to other objects in the graph.

Original: The graph can be thought as a succinct latent representation that ignores irrelevant object information in text observations when predicting the next state z t .
Revised: The graph can be thought as a succinct latent representation that ignores irrelevant object information in text observations when predicting the next state z t .

Original: (c) Independent Transition Layers.
Revised: (c) Independent Transition Layers.

Original: TBG environments typically contain a large population of objects, and each action can only affect a few underlying objects.
Revised: TBG environments typically contain a large population of objects, and each action can only affect a few underlying objects.

Original: In order to improve the robustness of the dynamics, we utilize 1) a Bi-Directional Attention Flow (BIDAF) network (Seo et al., 2017) to identify the affected objects and 2) a group of transition layers to predict the belief state of objects by following the Independent Causal Mechanism (ICM) framework (Pearl, 2009).
Revised: In order to improve the robustness of the dynamics, we utilize 1) a Bi-Directional Attention Flow (BIDAF) network (Seo et al., 2017) to identify the affected objects and 2) a group of transition layers to predict the belief state of objects by following the Independent Causal Mechanism (ICM) framework (Pearl, 2009).

Original: Given e t − 1 and an action a t − 1 =
Revised: Given e t − 1 and an action a t − 1 =

Original: [ a 1 ,t − 1 , ..., a J,t − 1 ] (a sentence of J words), the BIDAF network learns an attention matrix B ae ∈ [0 , 1] K × J , where b k,j measures the similarity between the representation of the k -th object and the embedding of the j -th word.
Revised: [ a 1 ,t − 1 , ..., a J,t − 1 ] (a sentence of J words), the BIDAF network learns an attention matrix B ae ∈ [0 , 1] K × J , where b k,j measures the similarity between the representation of the k -th object and the embedding of the j -th word.

Original: To learn an action-aware representation of object , we compute 1) action-to-object attentions b e that indicate which object representations have the closest similarity to one of the action words and are hence critical for modelling the impact of the action words; 2) object-to-action attentions b ak that indicate which action words are most relevant to the k th object representation.
Revised: To learn an action-aware representation of object , we compute 1) action-to-object attentions b e that indicate which object representations have the closest similarity to one of the action words and are hence critical for modelling the impact of the action words; 2) object-to-action attentions b ak that indicate which action words are most relevant to the k th object representation.

Original: The action-aware object representation concatenates the attended action vector ν ak,t − 1 , the attended object representation ν et − 1 and the initial object representation e k,t − 1 : 
Revised: The action-aware object representation concatenates the attended action vector ν ak,t − 1 , the attended object representation ν et − 1 and the initial object representation e k,t − 1 : ν

Original: [Equation]
Revised: [Equation]

Original: t
Revised: b

Original: [Equation]
Revised: [Equation]

Original:  where b ak = softmax ( B aek, : ) ∈ [0 , 1] J , 
Revised: ,t [Equation] [Equation] where b ak = softmax ( B aek, : ) ∈ [0 , 1] J , ν

Original: [Equation]
Revised: [Equation]

Original: where b e = softmax (max col ( B ae )) ∈
Revised: where b e = softmax (max col ( B ae )) ∈

Original: [0 , 1] K , p k T ( z k,t | a t − 1 , z t − 1 ) =
Revised: [0 , 1] K , p k T ( z k,t | a t − 1 , z t − 1 ) =

Original: N ( µ k,t , σ k,t ) where [ µ k,t , σ k,t ] = ψ pk ([ ν ak,t − 1 , ν et − 1 , e k,t − 1 ]) (3)
Revised: N ( µ k,t , σ k,t ) where [ µ k,t , σ k,t ] = ψ pk ([ ν ak,t − 1 , ν et − 1 , e k,t − 1 ]) (3)

Original: Here ψ a is a transformer-based text embedding function (Vaswani et al., 2017) for the input action, and ψ pk is implemented by a MLP.
Revised: Here ψ a is a transformer-based text embedding function (Vaswani et al., 2017) for the input action, and ψ pk is implemented by a MLP.

Original: The transition layers ψ p 1 , . .
Revised: The transition layers ψ p 1 , . .

Original: . , ψ pK are parameterized independently.
Revised: . , ψ pK are parameterized independently.

Original: 3. O BJECT O RIENTED R EWARD M ODEL
Revised: 3. O BJECT O RIENTED R EWARD M ODEL

Original: We learn a reward model p r ( r t | z t , g t ) based on the predicted belief states.
Revised: We learn a reward model p r ( r t | z t , g t ) based on the predicted belief states.

Original: This reward model is important for learning goal-conditioned policies that can utilize the knowledge learned from a task to solve similar tasks, which accelerates learning and facilitates knowledge generalization.
Revised: This reward model is important for learning goal-conditioned policies that can utilize the knowledge learned from a task to solve similar tasks, which accelerates learning and facilitates knowledge generalization.

Original: To implement the reward model, we build a goal extractor (Appendix A.3) that extracts goals g t from object states z t =
Revised: To implement the reward model, we build a goal extractor (Appendix A.3) that extracts goals g t from object states z t =

Original: [ z 1 ,t , . .
Revised: [ z 1 ,t , . .

Original: , z K,t ] .
Revised: , z K,t ] .

Original: The goal is described by a sequence of words and symbols that follow the triplet format, e.g., "[potato, -, need, -, fry]" where "-" serves as separator between object and relation mentions.
Revised: The goal is described by a sequence of words and symbols that follow the triplet format, e.g., "[potato, -, need, -, fry]" where "-" serves as separator between object and relation mentions.

Original: Although z t implicitly embeds g t with latent values, we prefer to represent g t explicitly with a sentence for the ease of model comprehension.
Revised: Although z t implicitly embeds g t with latent values, we prefer to represent g t explicitly with a sentence for the ease of model comprehension.

Original: Given the goal g t and object states z t , the reward model p r is deﬁned by
Revised: Given the goal g t and object states z t , the reward model p r is deﬁned by

Original: [Equation]
Revised: [Equation]

Original: where z[Equation] = [Equation]
Revised: z

Original: [Equation]
Revised: [Equation]

Original: ψ g is a transformer-based text embedding function for goals, and ψ r is implemented by a MLP.
Revised: ψ g is a transformer-based text embedding function for goals, and ψ r is implemented by a MLP.

Original: 3. T RAINING O BJECTIVE FOR D YNAMICS M ODELS
Revised: 3. T RAINING O BJECTIVE FOR D YNAMICS M ODELS

Original: In this section, we introduce the objectives for learning the transition model and the reward model under the object-supervised setting and the self-supervised setting.
Revised: In this section, we introduce the objectives for learning the transition model and the reward model under the object-supervised setting and the self-supervised setting.

Original: Object-Supervised Setting.
Revised: Object-Supervised Setting.

Original: We learn the dynamics model under the supervision of memory graphs h t that record rich information about the observed objects and their relations from the beginning of a game (Section 2.3).
Revised: We learn the dynamics model under the supervision of memory graphs h t that record rich information about the observed objects and their relations from the beginning of a game (Section 2.3).

Original: The Object-Supervised Evidence Lower Bound (OS-ELBo) objective is:
Revised: The Object-Supervised Evidence Lower Bound (OS-ELBo) objective is:

Original: T (cid:88) t
Revised: T (cid:88) t

Original: = (cid:110)E q
Revised: = (cid:110)

Original: [Equation] log [Equation] −
Revised: E q E (cid:104) log p Ω (ˆ h t | z t ) (cid:105) −

Original: K (cid:88) k
Revised: K (cid:88) k

Original: =
Revised: =

Original: D kl (cid:104)

Original: [Equation]

Original: k

Original: [Equation]

Original: o

Original: [Equation]

Original: a

Original: [Equation]

Original: [Equation]

Original: z

Original: [Equation]

Original: ) (cid:107)

Original: [Equation]

Original: ( z k,

Original: [Equation]

Original:  |  a [Equation] [Equation] ) (cid:105)(cid:111) (5) where p Ω is the graph decoder, p T is the transition model and q k E is the observation encoder.
Revised: D kl (cid:104) q k E ( z k,t | o t , a t − 1 , z t − 1 ) (cid:107) p k T ( z k,t | a t − 1 , z t − 1 ) (cid:105)(cid:111) (5) where p Ω is the graph decoder, p T is the transition model and q k E is the observation encoder.

Original: We now provide more details about each component and how to obtain h t .
Revised: We now provide more details about each component and how to obtain h t .

Original: p Ω is approximated by a re-sampling technique described in Appendix A.2.
Revised: p Ω is approximated by a re-sampling technique described in Appendix A.2.

Original: p T is a Gaussian distribution as described in
Revised: p T is a Gaussian distribution as described in

Original: Equation 3.
Revised: Equation 3.

Original: q k E shares a similar structure to p T (see Appendix A.5 for more details).
Revised: q k E shares a similar structure to p T (see Appendix A.5 for more details).

Original: The memory graph h t is not a direct output from the RL environment.
Revised: The memory graph h t is not a direct output from the RL environment.

Original: To perform object-supervised training, we extract object information from o t and predict ˆ h t with a deterministic object extractor f e ( o t , a t − 1 , h t − 1 ) .
Revised: To perform object-supervised training, we extract object information from o t and predict ˆ h t with a deterministic object extractor f e ( o t , a t − 1 , h t − 1 ) .

Original: We learn f e during a pre-training stage by utilizing a FTWP dataset.
Revised: We learn f e during a pre-training stage by utilizing a FTWP dataset.

Original: The dataset records the trajectories [ o 1 , a 1 , . . .
Revised: The dataset records the trajectories [ o 1 , a 1 , . . .

Original: , o t , a t ] and corresponding [ h 1 , . . .
Revised: , o t , a t ] and corresponding [ h 1 , . . .

Original: , h t ] by following walkthroughs in FTWP games.
Revised: , h t ] by following walkthroughs in FTWP games.

Original: In practice, we can build such a dataset by extracting triplets ( φ ) from the text observations ( o 1 , . . .
Revised: In practice, we can build such a dataset by extracting triplets ( φ ) from the text observations ( o 1 , . . .

Original: , o t ) with an Open Information Extractor (OIE) (Angeli et al., 2015) and map the captured triplets into a knowledge graph by following (Ammanabrolu & Hausknecht, 2020).
Revised: , o t ) with an Open Information Extractor (OIE) (Angeli et al., 2015) and map the captured triplets into a knowledge graph by following (Ammanabrolu & Hausknecht, 2020).

Original: After training f e , we initialize h 

Revised: After training f e , we initialize h

Original: = { 0 } C × K × K and update by ˆ h t =
Revised: = { 0 } C × K × K and update by ˆ h t =

Original: f e ( o t , a t − 1 , ˆ h t − 1 ) during training .
Revised: f e ( o t , a t − 1 , ˆ h t − 1 ) during training .

Original: The detailed introduction of f e and the FTWP dataset can be found in Appendix A.4 and A.6.
Revised: The detailed introduction of f e and the FTWP dataset can be found in Appendix A.4 and A.6.

Original: To train the reward model, we sample object states z t =
Revised: To train the reward model, we sample object states z t =

Original: [ z 1 ,t , . .
Revised: [ z 1 ,t , . .

Original: , z K,t ] from their belief states, extract goal g t , and minimize a smooth L1 loss (Girshick, 2015):
Revised: , z K,t ] from their belief states, extract goal g t , and minimize a smooth L1 loss (Girshick, 2015):

Original: [Equation]
Revised: [Equation]

Original: Self-Supervised Setting.
Revised: Self-Supervised Setting.

Original: In order to generalize our dynamics model to common RL environments where the predicted or ground truth memory graphs are unavailable, we develop a Self-Supervised
Revised: In order to generalize our dynamics model to common RL environments where the predicted or ground truth memory graphs are unavailable, we develop a Self-Supervised

Original: Evidence Lower Bound (SS-ELBo) objective that directly learns from rewards and observations:
Revised: Evidence Lower Bound (SS-ELBo) objective that directly learns from rewards and observations:

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

Original: where p T , p r and q E are the aforementioned transition model, reward model and observation encoder.
Revised: where p T , p r and q E are the aforementioned transition model, reward model and observation encoder.

Original: The observation decoder p o is implemented by a Sequence-to-Sequence (Seq2Seq) model (Sutskever et al., 2014), which is trained by a teacher-forcing technique with text observations o t .
Revised: The observation decoder p o is implemented by a Sequence-to-Sequence (Seq2Seq) model (Sutskever et al., 2014), which is trained by a teacher-forcing technique with text observations o t .

Original: This SS-EBLo objective enables training the reward and transition model together by maximizing one objective, which improves the training efﬁciency.
Revised: This SS-EBLo objective enables training the reward and transition model together by maximizing one objective, which improves the training efﬁciency.

Original: Appendix C.1 shows an example of a predicted graph.
Revised: Appendix C.1 shows an example of a predicted graph.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Environments.
Revised: Environments.

Original: We divide the games in the Text-World benchmark (Côté et al., 2018) into ﬁve subsets according to their difﬁculty levels.
Revised: We divide the games in the Text-World benchmark (Côté et al., 2018) into ﬁve subsets according to their difﬁculty levels.

Original: Each subset contains 100 training, 20 validation, and 20 testing games.
Revised: Each subset contains 100 training, 20 validation, and 20 testing games.

Original: These subsets are mutually exclusive.
Revised: These subsets are mutually exclusive.

Original: For the easier cooking games, the recipe requires a single ingredient and the world is limited to a single location, whereas harder games require an agent to navigate a map of multiple locations to collect and appropriately process up to three ingredients.
Revised: For the easier cooking games, the recipe requires a single ingredient and the world is limited to a single location, whereas harder games require an agent to navigate a map of multiple locations to collect and appropriately process up to three ingredients.

Original: Table 4 summarizes the game statistics.
Revised: Table 4 summarizes the game statistics.

Original: Running Settings.
Revised: Running Settings.

Original: To efﬁciently collect samples from the games in the training dataset , we train our OOTD model by implementing the Dyna-Q algorithm (Algorithm 1 in Appendix) due to the fact that 1) Q-network supports batch input, which enables solving and gathering data from multiple games in parallel, and 2) the Q-based algorithms provide a fair comparison of sample efﬁciency with other model-free baselines.
Revised: To efﬁciently collect samples from the games in the training dataset , we train our OOTD model by implementing the Dyna-Q algorithm (Algorithm 1 in Appendix) due to the fact that 1) Q-network supports batch input, which enables solving and gathering data from multiple games in parallel, and 2) the Q-based algorithms provide a fair comparison of sample efﬁciency with other model-free baselines.

Original: After training, we ﬁx the parameters in dynamics models and predict belief states and rewards to solve the games in the test dataset .
Revised: After training, we ﬁx the parameters in dynamics models and predict belief states and rewards to solve the games in the test dataset .

Original: During testing , we study the options of applying different planning algorithms including Dyna-Q, MCTS, and their combination (Dyna-Q+MCTS) that initializes MCTS with values from Q networks (Algorithm 2 in appendix).
Revised: During testing , we study the options of applying different planning algorithms including Dyna-Q, MCTS, and their combination (Dyna-Q+MCTS) that initializes MCTS with values from Q networks (Algorithm 2 in appendix).

Original: The model hyper-parameters are tuned with the games in the validation set.
Revised: The model hyper-parameters are tuned with the games in the validation set.

Original: (Appendix A.1).
Revised: (Appendix A.1).

Original: Baselines.
Revised: Baselines.

Original: we compare the state-of-the-art models for solving TBGs.
Revised: we compare the state-of-the-art models for solving TBGs.

Original: The baselines include Deep
Revised: The baselines include Deep

Original: Q-Network ( DQN ) (Narasimhan et al., 2015) and Deep Recurrent Q-Network ( DRQN )
Revised: Q-Network ( DQN ) (Narasimhan et al., 2015) and Deep Recurrent Q-Network ( DRQN )

Original: (Yuan et al., 2018).
Revised: (Yuan et al., 2018).

Original: DQN selects the action based on the current observation, whereas DRQN models the history of action and observations with an RNN.
Revised: DQN selects the action based on the current observation, whereas DRQN models the history of action and observations with an RNN.

Original: For fair comparisons, we replace their LSTM-based text encoders with transformers.
Revised: For fair comparisons, we replace their LSTM-based text encoders with transformers.

Original: We compare an extension of DRQN by using an episodic counting bonus to encourage exploration, which is denoted as ( DRQN+ ) (Yuan et al., 2018).
Revised: We compare an extension of DRQN by using an episodic counting bonus to encourage exploration, which is denoted as ( DRQN+ ) (Yuan et al., 2018).

Original: Our next baseline is KGA2C (Ammanabrolu & Hausknecht, 2020), which builds a graph-based state space and applies the
Revised: Our next baseline is KGA2C (Ammanabrolu & Hausknecht, 2020), which builds a graph-based state space and applies the

Original: Actor-Critic algorithm for learning the policy.
Revised: Actor-Critic algorithm for learning the policy.

Original: We replace their action generator with an action selector for comparison.
Revised: We replace their action generator with an action selector for comparison.

Original: The last baseline is GATA (Adhikari et al., 2020) that learns a graph-structured representation to model the environment dynamics and select actions based on this representation.
Revised: The last baseline is GATA (Adhikari et al., 2020) that learns a graph-structured representation to model the environment dynamics and select actions based on this representation.

Original: We experiment with several options of GATA, including: 1) GATA-GTP that pre-trains the discrete graph updater with ground-truth graphs from the FTWP dataset (Appendix A.6).
Revised: We experiment with several options of GATA, including: 1) GATA-GTP that pre-trains the discrete graph updater with ground-truth graphs from the FTWP dataset (Appendix A.6).

Original: Note that we apply the same pre-training dataset for our OS-ELBo objective.
Revised: Note that we apply the same pre-training dataset for our OS-ELBo objective.

Original: 2) GATA-OG that learns a continuous graph updater by reconstructing the text observation.
Revised: 2) GATA-OG that learns a continuous graph updater by reconstructing the text observation.

Original: GATA-COC that trains the graph updater with a contrastive observation classiﬁcation objective.
Revised: GATA-COC that trains the graph updater with a contrastive observation classiﬁcation objective.

Original: Both GATA-OG and GATA-COC follow the self-supervised setting, which offers fair comparison to our SS-ELBo objective.
Revised: Both GATA-OG and GATA-COC follow the self-supervised setting, which offers fair comparison to our SS-ELBo objective.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: C ONTROLLING P ERFORMANCE IN TBG E NVIRONMENTS
Revised: C ONTROLLING P ERFORMANCE IN TBG E NVIRONMENTS

Original: Table 2 shows the agent’s controlling performance at different difﬁculty levels.
Revised: Table 2 shows the agent’s controlling performance at different difﬁculty levels.

Original: Equipped with the
Revised: Equipped with the

Original: OOTD model, our agent implements model-based planning and achieves leading performance over the agents based on model-free algorithms.
Revised: OOTD model, our agent implements model-based planning and achieves leading performance over the agents based on model-free algorithms.

Original: We ﬁnd the SS-based OOTD models outperform the
Revised: We ﬁnd the SS-based OOTD models outperform the

Original: OS-based models by scoring more rewards at 4 (out of 6) difﬁculty levels.
Revised: OS-based models by scoring more rewards at 4 (out of 6) difﬁculty levels.

Original: Although the object signals facilitate a better prediction of dynamics, this advantage can not be directly transferred to controlling performance.
Revised: Although the object signals facilitate a better prediction of dynamics, this advantage can not be directly transferred to controlling performance.

Original: This is commonly known as the objective-mismatch problem in modelbased RL (Lambert et al., 2020).
Revised: This is commonly known as the objective-mismatch problem in modelbased RL (Lambert et al., 2020).

Original: On the other hand, SS-based models are directly optimized toward generating a better prediction of rewards, which are important signals for controlling.
Revised: On the other hand, SS-based models are directly optimized toward generating a better prediction of rewards, which are important signals for controlling.

Original: The rewardirrelevant information is ﬁltered to derive a more compressed representation for better generalization ability.
Revised: The rewardirrelevant information is ﬁltered to derive a more compressed representation for better generalization ability.

Original: Another important ﬁnding is that our OOTD-based agents perform better than other graphbased agents (e.g., GATA-GTP and KG-A2C).
Revised: Another important ﬁnding is that our OOTD-based agents perform better than other graphbased agents (e.g., GATA-GTP and KG-A2C).

Original: This is because OOTD is doing model-based RL by learning the dynamics at the object level while other graph-based techniques are model free since they learn to extract object representations, but not to predict their dynamics.
Revised: This is because OOTD is doing model-based RL by learning the dynamics at the object level while other graph-based techniques are model free since they learn to extract object representations, but not to predict their dynamics.

Original: Sample Efﬁciency.
Revised: Sample Efﬁciency.

Original: We empirically demonstrate the sample efﬁciency of our model-based agents by comparing their training performance with other model-free agents.
Revised: We empirically demonstrate the sample efﬁciency of our model-based agents by comparing their training performance with other model-free agents.

Original: This experiment studies three variations of proposed methods, including 1) OS-Dyna-Q and 2) SS-Dyna-Q that implement the Q-Dyna algorithm and train the OOTD models with the OS-ELBo objective and the SS-ELBo objective respectively.
Revised: This experiment studies three variations of proposed methods, including 1) OS-Dyna-Q and 2) SS-Dyna-Q that implement the Q-Dyna algorithm and train the OOTD models with the OS-ELBo objective and the SS-ELBo objective respectively.

Original: 3) Object-Oriented (OO)-DQN that trains the OOTD model to predict object states, based on which we compute action values.
Revised: 3) Object-Oriented (OO)-DQN that trains the OOTD model to predict object states, based on which we compute action values.

Original: As a model-free method, OO-DQN does not expand the replay buffer with the samples generated by dynamics models.
Revised: As a model-free method, OO-DQN does not expand the replay buffer with the samples generated by dynamics models.

Original: Figure 3 shows the training curve for 6 difﬁculty levels.
Revised: Figure 3 shows the training curve for 6 difﬁculty levels.

Original: The learning speed of proposed model-based algorithms (OS-Dyna-Q and SS-Dyna-Q) is generally faster than that of MF algorithms, but we observe an exception at the games of difﬁculty level 5, where GATA-OG converges faster.
Revised: The learning speed of proposed model-based algorithms (OS-Dyna-Q and SS-Dyna-Q) is generally faster than that of MF algorithms, but we observe an exception at the games of difﬁculty level 5, where GATA-OG converges faster.

Original: We ﬁnd it generally takes more samples to learn the interactions among K objects and their relations to action values.
Revised: We ﬁnd it generally takes more samples to learn the interactions among K objects and their relations to action values.

Original: It explains why (MF)-Dyna-Q converges slower than others.
Revised: It explains why (MF)-Dyna-Q converges slower than others.

Original: However, the object-based value function performs better than that based on raw observation in terms of normalized scores.
Revised: However, the object-based value function performs better than that based on raw observation in terms of normalized scores.

Original: 4. A BLATION S TUDIES FOR O BJECT -O RIENTED D YNAMICS
Revised: 4. A BLATION S TUDIES FOR O BJECT -O RIENTED D YNAMICS

Original: To understand the importance of our OOTD design, we empirically study the options of 1) removing the BIDAF attentions by directly applying the action and node embeddings to predict object states (i.e.,
Revised: To understand the importance of our OOTD design, we empirically study the options of 1) removing the BIDAF attentions by directly applying the action and node embeddings to predict object states (i.e.,

Original: No-Attentions Dynamics ( NAt-Dyna )) and 2) removing independent transition layers by building a single layer to predict states for all objects (i.e., Single-Layer Dynamics ( SLa-Dyna )).
Revised: No-Attentions Dynamics ( NAt-Dyna )) and 2) removing independent transition layers by building a single layer to predict states for all objects (i.e., Single-Layer Dynamics ( SLa-Dyna )).

Original: We compare these models with our OOTD models trained by the OS-ELBo objective ( OS-OOTD ) and the SSELBo objective ( SS-OOTD ).
Revised: We compare these models with our OOTD models trained by the OS-ELBo objective ( OS-OOTD ) and the SSELBo objective ( SS-OOTD ).

Original: To evaluate how well the modiﬁed model p k T (cid:48) captures object dynamics, we sample object states z k,t ∼ p k T (cid:48)
Revised: To evaluate how well the modiﬁed model p k T (cid:48) captures object dynamics, we sample object states z k,t ∼ p k T (cid:48)

Original: ( z k,t | a t − 1 , z t − 1 ) for a total of K objects, visualize these states with the t-distributed Stochastic Neighbor Embedding (t-SNE)
Revised: ( z k,t | a t − 1 , z t − 1 ) for a total of K objects, visualize these states with the t-distributed Stochastic Neighbor Embedding (t-SNE)

Original: algorithm (Van der Maaten & Hinton, 2008), and quantify the dynamics prediction performance.
Revised: algorithm (Van der Maaten & Hinton, 2008), and quantify the dynamics prediction performance.

Original: Object States Visualization.
Revised: Object States Visualization.

Original: In this experiment, we randomly pick 20 games from the testing games of all difﬁculty levels.
Revised: In this experiment, we randomly pick 20 games from the testing games of all difﬁculty levels.

Original: Within each game, p k T (cid:48) computes the states z k, 1 , · · · , z k,T with randomly selected actions a 0 , . .
Revised: Within each game, p k T (cid:48) computes the states z k, 1 , · · · , z k,T with randomly selected actions a 0 , . .

Original: . , a T − 1 , and t-SNE embeds z k,t into a vector of 2 dimensions.
Revised: . , a T − 1 , and t-SNE embeds z k,t into a vector of 2 dimensions.

Original: Figure 4 visualizes the embedded vectors labelled by 10 randomly selected objects.
Revised: Figure 4 visualizes the embedded vectors labelled by 10 randomly selected objects.

Original: The object states computed by our OOTD model are distinguishable, and thus t-SNE embeds the states for the same object close to each other.
Revised: The object states computed by our OOTD model are distinguishable, and thus t-SNE embeds the states for the same object close to each other.

Original: This phenomenon is not observable when we apply a single layer for computing object states (i.e., SLa-Dyna).
Revised: This phenomenon is not observable when we apply a single layer for computing object states (i.e., SLa-Dyna).

Original: Another important observation is that the OOTD model automatically learns to generate similar states for similar objects.
Revised: Another important observation is that the OOTD model automatically learns to generate similar states for similar objects.

Original: For example, OS-OOTD generates similar states for the object of ingredients (e.g., parsley, apple, carrot, and potato marked by the dark frame).
Revised: For example, OS-OOTD generates similar states for the object of ingredients (e.g., parsley, apple, carrot, and potato marked by the dark frame).

Original: A similar phenomenon, although less obvious, can be observed for SS-OOTD.
Revised: A similar phenomenon, although less obvious, can be observed for SS-OOTD.

Original: However, such a shrinkage effect disappears in NAt-Dyna, which explains the importance of bidirectional attentions.
Revised: However, such a shrinkage effect disappears in NAt-Dyna, which explains the importance of bidirectional attentions.

Original: Probing Experiment.
Revised: Probing Experiment.

Original: This experiment studies how well our OOTD models can capture useful information from the TBG environments.
Revised: This experiment studies how well our OOTD models can capture useful information from the TBG environments.

Original: We map the sampled object states z 1 ,t , . . .
Revised: We map the sampled object states z 1 ,t , . . .

Original: , z K,t to memory graphs ˆ h t and rewards ˆ r t with the additional graph decoder p (cid:48) G and reward model p (cid:48) r .
Revised: , z K,t to memory graphs ˆ h t and rewards ˆ r t with the additional graph decoder p (cid:48) G and reward model p (cid:48) r .

Original: The performance of the dynamics models is quantiﬁed by whether ˆ h t and ˆ r t can match the ground-truth h ∗ t and r ∗ t in the FTWP dataset (Appendix A.6).
Revised: The performance of the dynamics models is quantiﬁed by whether ˆ h t and ˆ r t can match the ground-truth h ∗ t and r ∗ t in the FTWP dataset (Appendix A.6).

Original: During sampling, the input actions a 1 , . . .
Revised: During sampling, the input actions a 1 , . . .

Original: , a T follow the recorded trajectories in FTWP.
Revised: , a T follow the recorded trajectories in FTWP.

Original: The ground-truth rewards in FTWP games is either one or zero, and we perform an entry-level evaluation (e.g., h t [ i, j, c ] ) on the adjacency matrix of a predicted graph.
Revised: The ground-truth rewards in FTWP games is either one or zero, and we perform an entry-level evaluation (e.g., h t [ i, j, c ] ) on the adjacency matrix of a predicted graph.

Original: It allows us to evaluate whether the real rewards and triplets are captured with F1 scores.
Revised: It allows us to evaluate whether the real rewards and triplets are captured with F1 scores.

Original: For a fair comparison, the games in the FTWP dataset are divided into training (80%), validation (10%) and testing (10%) sets, based on which we ﬁrst freeze the parameters in the dynamics models and then update the added p (cid:48)G and p (cid:48) r with data in training games.
Revised: For a fair comparison, the games in the FTWP dataset are divided into training (80%), validation (10%) and testing (10%) sets, based on which we ﬁrst freeze the parameters in the dynamics models and then update the added p (cid:48)G and p (cid:48) r with data in training games.

Original: [Table]
Revised: [Table]

Original: Table 3 shows the testing performance.
Revised: Table 3 shows the testing performance.

Original: We study both 1) the graph- Generation performance which measures the accuracy of generating all the candidate triplets in a graph and 2) the graph- Updating performance which evaluates only the triplets that have been updated at the current step t .
Revised: We study both 1) the graph- Generation performance which measures the accuracy of generating all the candidate triplets in a graph and 2) the graph- Updating performance which evaluates only the triplets that have been updated at the current step t .

Original: We add a Real and a Random baseline that feed the groundtruth graph h ∗ t and a randomly generated graph h RMt into the graph encoder of transition models.
Revised: We add a Real and a Random baseline that feed the groundtruth graph h ∗ t and a randomly generated graph h RMt into the graph encoder of transition models.

Original: They study the probing performance when ground-truth graphs are known and when the posterior collapses for the transition model (so p k T (cid:48) ( z k,t | a t − 1 , z t − 1 ) =
Revised: They study the probing performance when ground-truth graphs are known and when the posterior collapses for the transition model (so p k T (cid:48) ( z k,t | a t − 1 , z t − 1 ) =

Original: p k T (cid:48)
Revised: p k T (cid:48)

Original: ( z k,t ) ).
Revised: ( z k,t ) ).

Original: The results show that
Revised: The results show that

Original: OS-OOTD achieves a better graph prediction performance by applying the object-supervised objectives (OS-ELBo), but SSOOTD has higher reward prediction accuracy, since SS-ELBo includes an end-to-end training to predict future rewards, which ﬁlters the reward-irrelevant object information.
Revised: OS-OOTD achieves a better graph prediction performance by applying the object-supervised objectives (OS-ELBo), but SSOOTD has higher reward prediction accuracy, since SS-ELBo includes an end-to-end training to predict future rewards, which ﬁlters the reward-irrelevant object information.

Original: Another observation is that removing independent layers and directional attentions signiﬁcantly degrade the graph prediction performance.
Revised: Another observation is that removing independent layers and directional attentions signiﬁcantly degrade the graph prediction performance.

Original: It demonstrates that they are important components of our OOTD model.
Revised: It demonstrates that they are important components of our OOTD model.

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

Original: DRL Agent for Text-Based Games (TBGs).
Revised: DRL Agent for Text-Based Games (TBGs).

Original: TBGs (e.g., Textworld (Côté et al., 2018) and interactive ﬁctional games (Hausknecht et al., 2020)) are interactive simulations with text observations and actions.
Revised: TBGs (e.g., Textworld (Côté et al., 2018) and interactive ﬁctional games (Hausknecht et al., 2020)) are interactive simulations with text observations and actions.

Original: To solve TBGs, previous DRL agents typically applied advanced neural models for representing the state and action sentences, including recurrent models (Narasimhan et al., 2015; Jain et al., 2020; Zelinka, 2018; Madotto et al., 2020; Adolphs & Hofmann, 2020),
Revised: To solve TBGs, previous DRL agents typically applied advanced neural models for representing the state and action sentences, including recurrent models (Narasimhan et al., 2015; Jain et al., 2020; Zelinka, 2018; Madotto et al., 2020; Adolphs & Hofmann, 2020),

Original: Transformers (Adhikari et al., 2020; Ammanabrolu et al., 2020; Urbanek et al., 2019), relevance networks (He et al., 2016), and convolutional neural networks (Yin & May, 2019; Zahavy et al., 2018; Yin & May, 2019).
Revised: Transformers (Adhikari et al., 2020; Ammanabrolu et al., 2020; Urbanek et al., 2019), relevance networks (He et al., 2016), and convolutional neural networks (Yin & May, 2019; Zahavy et al., 2018; Yin & May, 2019).

Original: Since states in the TBGs environment are partially observable, many previous studies have formulated
Revised: Since states in the TBGs environment are partially observable, many previous studies have formulated

Original: TBGs as POMDPs and modelled the history of observations with a Deep Recurrent Q-Network (DRQN)
Revised: TBGs as POMDPs and modelled the history of observations with a Deep Recurrent Q-Network (DRQN)

Original: (Yuan et al., 2018) or a knowledge graph (Hausknecht et al., 2019; Ammanabrolu & Riedl, 2019; Ammanabrolu & Hausknecht, 2020; Adhikari et al., 2020; Ammanabrolu et al., 2020).
Revised: (Yuan et al., 2018) or a knowledge graph (Hausknecht et al., 2019; Ammanabrolu & Riedl, 2019; Ammanabrolu & Hausknecht, 2020; Adhikari et al., 2020; Ammanabrolu et al., 2020).

Original: Previous works are commonly based on model-free Q learning while we explore model-based planning.
Revised: Previous works are commonly based on model-free Q learning while we explore model-based planning.

Original: Object Oriented Reinforcement Learning.
Revised: Object Oriented Reinforcement Learning.

Original: Diuk et al.
Revised: Diuk et al.

Original: (2008) proposed modelling the transition and reward models at the object level.
Revised: (2008) proposed modelling the transition and reward models at the object level.

Original: Previous work (Mohan & Laird, 2011; Joshi et al., 2012; da Silva & Costa, 2018; Marom & Rosman, 2018; da Silva et al., 2019) explored training DRL agents with the object-level features.
Revised: Previous work (Mohan & Laird, 2011; Joshi et al., 2012; da Silva & Costa, 2018; Marom & Rosman, 2018; da Silva et al., 2019) explored training DRL agents with the object-level features.

Original: However, object features are not commonly available from the RL environments.
Revised: However, object features are not commonly available from the RL environments.

Original: Recent works have extracted object representations from states with optical ﬂow (Goel et al., 2018), variational models (Burgess et al., 2019) and structured convolutional layers (Kipf et al., 2020; Finn et al., 2016; Zhu et al., 2018; 2020).
Revised: Recent works have extracted object representations from states with optical ﬂow (Goel et al., 2018), variational models (Burgess et al., 2019) and structured convolutional layers (Kipf et al., 2020; Finn et al., 2016; Zhu et al., 2018; 2020).

Original: Previous works commonly studied image-based fully observable environments whereas we explore text-based partially observable environments.
Revised: Previous works commonly studied image-based fully observable environments whereas we explore text-based partially observable environments.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We proposed an OOTD model that enables planning in text-based environments.
Revised: We proposed an OOTD model that enables planning in text-based environments.

Original: The OOTD model learned a graph representation for capturing object dynamics and predicted the belief of object states with independent transition layers.
Revised: The OOTD model learned a graph representation for capturing object dynamics and predicted the belief of object states with independent transition layers.

Original: Empirical results demonstrate that OOTD-based agents outperform model-free baselines in terms of controlling performance and sample efﬁciency.
Revised: Empirical results demonstrate that OOTD-based agents outperform model-free baselines in terms of controlling performance and sample efﬁciency.

Original: We then provided an ablation study to show the importance of OOTD components.
Revised: We then provided an ablation study to show the importance of OOTD components.

Original: A promising future direction will be expanding our OOTDmodel to parser-based games where the agent must generate the commands character by character without knowing candidate actions.
Revised: A promising future direction will be expanding our OOTDmodel to parser-based games where the agent must generate the commands character by character without knowing candidate actions.

Original:  E THICS S TATEMENT
Revised:  E THICS S TATEMENT

Original: In this work, we study model-based planning in a simulated (TBG) environment.
Revised: In this work, we study model-based planning in a simulated (TBG) environment.

Original: Although the task itself has limited consequences for society, we take a broader view of our work and discuss its potential inﬂuence on future research that might have a considerable impact on our society.
Revised: Although the task itself has limited consequences for society, we take a broader view of our work and discuss its potential inﬂuence on future research that might have a considerable impact on our society.

Original: We believe developing an autonomous RL agent for TBGs serves as the stepping stone toward real-world applications where the machine cooperates or interacts with humans for some speciﬁc tasks.
Revised: We believe developing an autonomous RL agent for TBGs serves as the stepping stone toward real-world applications where the machine cooperates or interacts with humans for some speciﬁc tasks.

Original: A concrete example of these applications will be the auto-dialogue systems, where the machine must respond to human concerns by conveying accurate information and solving the problem with humans.
Revised: A concrete example of these applications will be the auto-dialogue systems, where the machine must respond to human concerns by conveying accurate information and solving the problem with humans.

Original: During this process, failing to communicate clearly, sensibly, or convincingly might also cause harm, for example, without sufﬁcient calibration, the machine could generate improper words or phrases that might potentially cause a bias toward some minorities or underrepresented groups.
Revised: During this process, failing to communicate clearly, sensibly, or convincingly might also cause harm, for example, without sufﬁcient calibration, the machine could generate improper words or phrases that might potentially cause a bias toward some minorities or underrepresented groups.

Original: How to eliminate such bias will be an important direction of future research.
Revised: How to eliminate such bias will be an important direction of future research.

Original:  R EPRODUCIBILITY S TATEMENT
Revised:  R EPRODUCIBILITY S TATEMENT

Original: This section introduces how we facilitate the reproducibility from the following perspectives:
Revised: This section introduces how we facilitate the reproducibility from the following perspectives:

Original: Model Implementation.
Revised: Model Implementation.

Original: Section 4 introduces our approach to dividing the training, validation, and testing datasets as well as the running settings based on the constructed datasets.
Revised: Section 4 introduces our approach to dividing the training, validation, and testing datasets as well as the running settings based on the constructed datasets.

Original: We report the hyper-parameters for training our OOTD model and the seeds for independent runs in Appendix A.1.
Revised: We report the hyper-parameters for training our OOTD model and the seeds for independent runs in Appendix A.1.

Original: The computing resources and running time are discussed in Appendix A.7.
Revised: The computing resources and running time are discussed in Appendix A.7.

Original: Code and Dataset.
Revised: Code and Dataset.

Original: We uploaded the code to the submission system.
Revised: We uploaded the code to the submission system.

Original: A public version will be added to GitHub once the anonymous period ends.
Revised: Please ﬁnd our implementation on Github 1 .

Original: We include all the required packages and environments settings to implement our model in the text document of the uploaded code.
Revised: We include all the required packages and environments settings to implement our model in the text document of the uploaded code.

Original: In terms of the dataset, we used public and self-generated datasets.
Revised: In terms of the dataset, we used public and self-generated datasets.

Original: For the public one, we provide an introduction in Appendix A.6 and the corresponding download link.
Revised: For the public one, we provide an introduction in Appendix A.6 and the corresponding download link.

Original: For the self-generated one, we have uploaded the dataset to our net drive.
Revised: For the self-generated one, we have uploaded the dataset to our net drive.

Original: The download link will be included in the uploaded code once the anonymous period ends.
Revised: The download link will be included in the uploaded code once the anonymous period ends.

Original: A M ODEL I MPLEMENTATION D ETAILS AND D ATASETS A. H YPER -P ARAMETERS
Revised: A M ODEL I MPLEMENTATION D ETAILS AND D ATASETS A. H YPER -P ARAMETERS

Original: We report the key hyper-parameters that have been applied for model implementation.
Revised: We report the key hyper-parameters that have been applied for model implementation.

Original: We set the number of candidate objects K to 99 and the number of candidate relations to 10 by following the settings in (Adhikari et al., 2020), which show that the agent can properly solve a game by modelling these objects and relations.
Revised: We set the number of candidate objects K to 99 and the number of candidate relations to 10 by following the settings in (Adhikari et al., 2020), which show that the agent can properly solve a game by modelling these objects and relations.

Original: The learning rate of policy training and dynamics model training is set to 0.0001.
Revised: The learning rate of policy training and dynamics model training is set to 0.0001.

Original: The discount factor γ is set to 0.9.
Revised: The discount factor γ is set to 0.9.

Original: The size of the hidden layers and the size of z in our OOTD model are set to 32.
Revised: The size of the hidden layers and the size of z in our OOTD model are set to 32.

Original: The transition layers of different objects are parameterized independently, so the neural network predicts 32 × 99 = 3 , 168 values at each step during dynamics modelling.
Revised: The transition layers of different objects are parameterized independently, so the neural network predicts 32 × 99 = 3 , 168 values at each step during dynamics modelling.

Original: The sizes of word embedding and node embedding are set to 300 and 100 respectively.
Revised: The sizes of word embedding and node embedding are set to 300 and 100 respectively.

Original: These numbers are determined experimentally, which could provide promising performance with a reasonable computational resource.
Revised: These numbers are determined experimentally, which could provide promising performance with a reasonable computational resource.

Original: The full list of hyper-parameters is recorded in the .yaml ﬁle in our submitted code.
Revised: The full list of hyper-parameters is recorded in the .yaml ﬁle in our submitted code.

Original: In our experiment, the random seeds we select are 123, 321, and 666.
Revised: In our experiment, the random seeds we select are 123, 321, and 666.

Original: A. R ESAMPLING FOR C OMPUTING THE G RAPH L IKELIHOOD
Revised: A. R ESAMPLING FOR C OMPUTING THE G RAPH L IKELIHOOD

Original: The graph likelihood p Ω ( h t | z t ) computes the probability of generating h t with a graph decoder.
Revised: The graph likelihood p Ω ( h t | z t ) computes the probability of generating h t with a graph decoder.

Original: As it is discussed in Section 3.1, our graph decoder is implemented as a scoring function for entry-level prediction.
Revised: As it is discussed in Section 3.1, our graph decoder is implemented as a scoring function for entry-level prediction.

Original: However, h t is a graph with signiﬁcant sparsity, for example, the corresponding adjacency matrix assigns positive labels (i.e., 1) for an average of 97.6 out of all 98,010 entries.
Revised: However, h t is a graph with signiﬁcant sparsity, for example, the corresponding adjacency matrix assigns positive labels (i.e., 1) for an average of 97.6 out of all 98,010 entries.

Original: As a result, directly optimizing toward h ∗ t creates a large negative bias.
Revised: As a result, directly optimizing toward h ∗ t creates a large negative bias.

Original: In other words, the decoder can achieve a satisfying accuracy by labelling all the entries with 0s, which does not capture any object information.
Revised: In other words, the decoder can achieve a satisfying accuracy by labelling all the entries with 0s, which does not capture any object information.

Original: To solve this problem, we implement a re-sampling technique that extracts all the positive samples (the entries labelled with one) from h ∗ t and sample an equivalent number of negative samples from the rest of the entries.
Revised: To solve this problem, we implement a re-sampling technique that extracts all the positive samples (the entries labelled with one) from h ∗ t and sample an equivalent number of negative samples from the rest of the entries.

Original: These positive and negative samples are used to estimate the graph likelihood.
Revised: These positive and negative samples are used to estimate the graph likelihood.

Original: A. G OAL E XTRACTOR
Revised: A. G OAL E XTRACTOR

Original: At a time step t , the goal extractor generates the goal g t from the object states z t by optionally conditioning on the previous goal g t − 1 and observed text o t .
Revised: At a time step t , the goal extractor generates the goal g t from the object states z t by optionally conditioning on the previous goal g t − 1 and observed text o t .

Original: The goal can be represented by a sequence of words and symbols that follows the triplet format.
Revised: The goal can be represented by a sequence of words and symbols that follows the triplet format.

Original: To generate this goal, we introduce two kinds of extractors that are built for the OOTD learned by the OS-ELBo and the SS-ELBo objectives.
Revised: To generate this goal, we introduce two kinds of extractors that are built for the OOTD learned by the OS-ELBo and the SS-ELBo objectives.

Original: Goal Extractor with the OS-ELBo objective.
Revised: Goal Extractor with the OS-ELBo objective.

Original: Under this setting, the predicted memory graph ˆ h OSt captures the sparsity of object observations and forms a direct prediction of the ground truth graph ˆ h ∗ t (check Figure 7).
Revised: Under this setting, the predicted memory graph ˆ h OSt captures the sparsity of object observations and forms a direct prediction of the ground truth graph ˆ h ∗ t (check Figure 7).

Original: We can build a rule-based goal extractor that maps z t to ˆ h OSt and checks whether h OSt captures certain relations (e.g., ’need’) between two objects.
Revised: We can build a rule-based goal extractor that maps z t to ˆ h OSt and checks whether h OSt captures certain relations (e.g., ’need’) between two objects.

Original: If there is such a relation, we directly output the sentence of corresponding triplets by following a format "[ object 1 $ relation $ object 2 ]" .
Revised: If there is such a relation, we directly output the sentence of corresponding triplets by following a format "[ object 1 $ relation $ object 2 ]" .

Original: Such a simple extractor can generalize well to training and testing games if ˆ h OSt manages to capture correct relations.
Revised: Such a simple extractor can generalize well to training and testing games if ˆ h OSt manages to capture correct relations.

Original: Goal Extractor with the SS-ELBo objective.
Revised: Goal Extractor with the SS-ELBo objective.

Original: The memory graph ˆ h OSt learned by the SS-ELBo objective captures a latent representation of ˆ h ∗ t .
Revised: The memory graph ˆ h OSt learned by the SS-ELBo objective captures a latent representation of ˆ h ∗ t .

Original: This latent representation is not directly interpretable, and thus we build a word-by-word sequence generator whose output is the goal to be predicted ( g t ) and inputs are z t , g t − 1 and o t .
Revised: This latent representation is not directly interpretable, and thus we build a word-by-word sequence generator whose output is the goal to be predicted ( g t ) and inputs are z t , g t − 1 and o t .

Original: To train this model, we manually build a goal learning dataset that collects pairs of model input/output from the training games.
Revised: To train this model, we manually build a goal learning dataset that collects pairs of model input/output from the training games.

Original: During testing, we predict the goal after interacting with the environment and update the current goal if a valid goal sentence is generated.
Revised: During testing, we predict the goal after interacting with the environment and update the current goal if a valid goal sentence is generated.

Original: A. O BJECT E XTRACTOR
Revised: A. O BJECT E XTRACTOR

Original: The object extractor f e ( o t , a t − 1 , h t − 1 ) predicts the memory graph h t by conditioning on current observations o t and action a t .
Revised: The object extractor f e ( o t , a t − 1 , h t − 1 ) predicts the memory graph h t by conditioning on current observations o t and action a t .

Original: Figure 5 shows the architecture of f e .
Revised: Figure 5 shows the architecture of f e .

Original: It shares a similar structure to our OOTD model, but the key differences are 1) f e is a deterministic model which only maximizes the likelihood of predicting h ∗ t during training.
Revised: It shares a similar structure to our OOTD model, but the key differences are 1) f e is a deterministic model which only maximizes the likelihood of predicting h ∗ t during training.

Original: 2) f e captures object information from o t , whereas o t is not available for our OOTD model, and it is the motivation why OOTD predicts the belief (distribution) of states while f e directly generates the deterministic states.
Revised: 2) f e captures object information from o t , whereas o t is not available for our OOTD model, and it is the motivation why OOTD predicts the belief (distribution) of states while f e directly generates the deterministic states.

Original: [Figure]
Revised: [Figure]

Original: A. O BSERVATION E NCODER .
Revised: A. O BSERVATION E NCODER .

Original: The observation encoder q E is deﬁned by: ν ak,t − 1 = (cid:88)
Revised: The observation encoder q E is deﬁned by: ν ak,t − 1 = (cid:88)

Original: j b ak,j ψ a ( a j,t − 1 ) where b ak = softmax ( B aek, : ) ∈ [0 , 1] J , ν et − 1 = (cid:88)
Revised: j b ak,j ψ a ( a j,t − 1 ) where b ak = softmax ( B aek, : ) ∈ [0 , 1] J , ν et − 1 = (cid:88)

Original: k b ek e k,t −where b e = softmax
Revised: k b ek e k,t −where b e = softmax

Original: (max col ( B ae )) ∈ [0 , 1] K , ν ok,t − 1 = (cid:88)
Revised: (max col ( B ae )) ∈ [0 , 1] K , ν ok,t − 1 = (cid:88)

Original: j b ok,j ψ o ( a j,t − 1 ) where b ok = softmax ( B oek, : ) ∈ [0 , 1] J ν (cid:48) ,et − 1 = (cid:88)
Revised: j b ok,j ψ o ( a j,t − 1 ) where b ok = softmax ( B oek, : ) ∈ [0 , 1] J ν (cid:48) ,et − 1 = (cid:88)

Original: k b oe,ek e k,t −where b oe,e = softmax (max col ( B oe )) ∈ [0 , 1] K (8) q k E ( z k,t
Revised: k b oe,ek e k,t −where b oe,e = softmax (max col ( B oe )) ∈ [0 , 1] K (8) q k E ( z k,t

Original: | o t , a t − 1 , z k,t − 1 ) =
Revised: | o t , a t − 1 , z k,t − 1 ) =

Original: N ( µ qk,t , σ qk,t ) where [ µ qk,t , σ qk,t ] = ψ qk ([ ν ok,t − 1 , ν (cid:48) ,et − 1 , ν ak,t − 1 , ν et − 1 , e k,t − 1 ]) ψ o is a transformer-based text embedding function (Vaswani et al., 2017) for observations, and ψ qk is implemented by a MLP.
Revised: N ( µ qk,t , σ qk,t ) where [ µ qk,t , σ qk,t ] = ψ qk ([ ν ok,t − 1 , ν (cid:48) ,et − 1 , ν ak,t − 1 , ν et − 1 , e k,t − 1 ]) ψ o is a transformer-based text embedding function (Vaswani et al., 2017) for observations, and ψ qk is implemented by a MLP.

Original: The transition layers ψ q 1 , . .
Revised: The transition layers ψ q 1 , . .

Original: . , ψ qK are parameterized independently.
Revised: . , ψ qK are parameterized independently.

Original: Figure 6 shows the structure of an observation encoder q k E , which shares a similar structure to the transition (OOTD) model (Figure 3), but it includes an additional BIDAF network to learn an attention matrix B oe , where b oek,i measures the similarity between the representation of the k th object and the i th word in the current observation o t .
Revised: Figure 6 shows the structure of an observation encoder q k E , which shares a similar structure to the transition (OOTD) model (Figure 3), but it includes an additional BIDAF network to learn an attention matrix B oe , where b oek,i measures the similarity between the representation of the k th object and the i th word in the current observation o t .

Original: It predicts the memory graph h t from z qt with a ComplEx graph decoder (Section 3.1).
Revised: It predicts the memory graph h t from z qt with a ComplEx graph decoder (Section 3.1).

Original: It predicts the current observation ˆ o t and reward r t with a sequence decoder and the reward function (Section 3.2).
Revised: It predicts the current observation ˆ o t and reward r t with a sequence decoder and the reward function (Section 3.2).

Original: [Figure]
Revised: [Figure]

Original: Graph/ sequenceDecoder
Revised: Graph/ sequenceDecoder

Original: [Figure]
Revised: [Figure]

Original:  the OS/SS ELBo
Revised: For the OS/SS ELBo

Original: For the SS ELBo
Revised: For the SS ELBo

Original: [Figure]
Revised: [Figure]

Original: A. FTWP DATASET
Revised: A. FTWP DATASET

Original: The FTWP dataset is a public dataset that supports pre-training the dynamics model 1 .
Revised: The FTWP dataset is a public dataset that supports pre-training the dynamics model 2 .

Original: Trischler et al.
Revised: Trischler et al.

Original: (2019) provided the First TextWorld Problems (FTWP) dataset.
Revised: (2019) provided the First TextWorld Problems (FTWP) dataset.

Original: The dataset consists of TextWorld games in a cooking theme across a wide range of difﬁculty levels.
Revised: The dataset consists of TextWorld games in a cooking theme across a wide range of difﬁculty levels.

Original: For each game, the expert trajectory and the memory graph h t at each time step are recorded.
Revised: For each game, the expert trajectory and the memory graph h t at each time step are recorded.

Original: The original dataset has only 10 games for each difﬁculty level.
Revised: The original dataset has only 10 games for each difﬁculty level.

Original: This is insufﬁcient for dynamics model pre-training, so Adhikari et al.
Revised: This is insufﬁcient for dynamics model pre-training, so Adhikari et al.

Original: (2020) further expanded this dataset by adding more games into this dataset by utilizing the logic in the
Revised: (2020) further expanded this dataset by adding more games into this dataset by utilizing the logic in the

Original: Textworld game engine.
Revised: Textworld game engine.

Original: The speciﬁc generation rules are described in (Adhikari et al., 2020).
Revised: The speciﬁc generation rules are described in (Adhikari et al., 2020).

Original: After expansion, the dataset contains 4440 games for training, 222 games for validation, 514 games for testing.
Revised: After expansion, the dataset contains 4440 games for training, 222 games for validation, 514 games for testing.

Original: To ensure fairness when using this dataset, we conﬁrm that there is no overlap between the
Revised: To ensure fairness when using this dataset, we conﬁrm that there is no overlap between the

Original: FTWP and the games we used to train and evaluate our planning algorithms.
Revised: FTWP and the games we used to train and evaluate our planning algorithms.

Original: A. C OMPUTATIONAL R ESOURCE AND R UNNING T IME
Revised: A. C OMPUTATIONAL R ESOURCE AND R UNNING T IME

Original: We run the experiment on a cluster operated by the Slurm workload manager.
Revised: We run the experiment on a cluster operated by the Slurm workload manager.

Original: The cluster has multiple kinds of GPUs, including Tesla T4 with 16 GB memory, Tesla P100 with 12 GB memory, and RTX 6000 with 24 GB memory.
Revised: The cluster has multiple kinds of GPUs, including Tesla T4 with 16 GB memory, Tesla P100 with 12 GB memory, and RTX 6000 with 24 GB memory.

Original: We used machines with 24 GB of memory for pre-training the object extractor and 64 GB for training the OOTD model.
Revised: We used machines with 24 GB of memory for pre-training the object extractor and 64 GB for training the OOTD model.

Original: The number of running nodes is 1, and the number of CPUs requested per task is 32.
Revised: The number of running nodes is 1, and the number of CPUs requested per task is 32.

Original: Given the aforementioned resources, pre-training can be completed within 4 days (96 hours), and training the dynamics model takes 6 days (144 hours).
Revised: Given the aforementioned resources, pre-training can be completed within 4 days (96 hours), and training the dynamics model takes 6 days (144 hours).

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: B A LGORITHMS
Revised: B A LGORITHMS

Revised: Algorithm

Original: Input : Transition Model: p T ( z t +1 | a t , z t ) ,
Revised: Input : Transition Model: p T ( z t +1 | a t , z t ) ,

Original: Object Encoder: q E ( z t +1 |
Revised: Object Encoder: q E ( z t +1 |

Original: o t +1 , a t , z t ) ,
Revised: o t +1 , a t , z t ) ,

Original: Reward Model: p r ( r t +1 | z t +1 , a t , g t ) ,
Revised: Reward Model: p r ( r t +1 | z t +1 , a t , g t ) ,

Original: Object Extraction function: f e ( h t , a t , o t +1 ) ,
Revised: Object Extraction function: f e ( h t , a t , o t +1 ) ,

Original: Graph Decoder: p Ω ( h t | z t ) ,
Revised: Graph Decoder: p Ω ( h t | z t ) ,

Original: Training game ids I .
Revised: Training game ids I .

Original: Initialize Replay Buffers D ; while not converged do
Revised: Initialize Replay Buffers D ; while not converged do

Original: Initialize belief state z 0 = random numbers, action a 0 = " restart ", goal g 0 = " <pad> " and memory graph h 

Revised: Initialize belief state z 0 = random numbers, action a 0 = " restart ", goal g 0 = " <pad> " and memory graph h

Original: =
Revised: =

Original: C × K × K ;
Revised: C × K × K ;

Original: Randomly pick a game id ∈ I , set t = 0 ; while not end do /* Q learning with environment */ o t +1 , r t +1 ← env ( a t , id ) ;
Revised: Randomly pick a game id ∈ I , set t = 0 ; while not end do /* Q learning with environment */ o t +1 , r t +1 ← env ( a t , id ) ;

Original: Extract object information h t
Revised: Extract object information h t

Original: +1 =
Revised: +1 =

Original: f ( a t , z t , o t +1 ) ;
Revised: f ( a t , z t , o t +1 ) ;

Original: Update object belief z t +1 ∼ q E ( z t +1 | o t +1 , a t , z t ) ;
Revised: Update object belief z t +1 ∼ q E ( z t +1 | o t +1 , a t , z t ) ;

Original: Extract goal: g t +1 =
Revised: Extract goal: g t +1 =

Original: ExtactGoal ( z t +1 , o t +1 , g t ) ;
Revised: ExtactGoal ( z t +1 , o t +1 , g t ) ;

Original: Action selection: ˆ a t +1 =
Revised: Action selection: ˆ a t +1 =

Original: EpsilonGreedy ( Q ( z t +1 , a t +1 , g t +1 )) ;
Revised: EpsilonGreedy ( Q ( z t +1 , a t +1 , g t +1 )) ;

Original: Expand the replay buffer: D ∪ { z t , h t , a t , g t , r t +1 , z t +1 , h t +1 , g t +1 } ; t = t + 1 and end =
Revised: Expand the replay buffer: D ∪ { z t , h t , a t , g t , r t +1 , z t +1 , h t +1 , g t +1 } ; t = t + 1 and end =

Original: CheckEnd ( env ) endfor update step c=1,...,C do /*
Revised: CheckEnd ( env ) endfor update step c=1,...,C do /*

Original: Model fitting */
Revised: Model fitting */

Original: Sample a transition: T = { z t , h t , a t , g t , r t +1 , z t +1 , h t +1 , g t +1 } from the dataset D ;
Revised: Sample a transition: T = { z t , h t , a t , g t , r t +1 , z t +1 , h t +1 , g t +1 } from the dataset D ;

Original: Update the reward model to minimize the L1Smooth loss (Equation 6) with T ;
Revised: Update the reward model to minimize the L1Smooth loss (Equation 6) with T ;

Original: Update the transition model to minimize the OS-ELBO (Equation 5) objective or the
Revised: Update the transition model to minimize the OS-ELBO (Equation 5) objective or the

Original: SS-ELBO objective (Equation 7) with T ;
Revised: SS-ELBO objective (Equation 7) with T ;

Original: Update the Q function to minimize the TD Loss (Equation 1) with T ; for explore step b=1,...,B do /* Exploration with the transition and reward models */
Revised: Update the Q function to minimize the TD Loss (Equation 1) with T ; for explore step b B do /* Exploration with the transition and reward models */

Original: Randomly select a (cid:48)
Revised: Randomly select a (cid:48)

Original: τ ∈ A ; z
Revised: τ ∈ A ; z

Original: [Equation]
Revised: [Equation]

Original: Update the transition T (cid:48)
Revised: Update the transition T (cid:48)

Original: = { z t , g t , a (cid:48) t , r (cid:48) t +1 , z (cid:48)
Revised: = { z t , g t , a (cid:48) t , r (cid:48) t +1 , z (cid:48)

Original: t +1 , g t +1 } ;
Revised: t +1 , g t +1 } ;

Original: Update the Q function to minimize the TD Loss (Equation 1) with T ; end end end
Revised: Update the Q function to minimize the TD Loss (Equation 1) with T ; end end end

Original: Input : Transition Model: p T ( z t +1 | a t , z t ) ,
Revised: Input : Transition Model: p T ( z t +1 | a t , z t ) ,

Original: Object Encoder: q E ( z t +1 |
Revised: Object Encoder: q E ( z t +1 |

Original: o t +1 , a t , z t ) ,
Revised: o t +1 , a t , z t ) ,

Original: Reward Model: p r ( r t +1 | z t +1 , a t , g t ) ,
Revised: Reward Model: p r ( r t +1 | z t +1 , a t , g t ) ,

Original: Action Value function: Q ( z t , a t , g t ) ,
Revised: Action Value function: Q ( z t , a t , g t ) ,

Original: Testing game ids
Revised: Testing game ids

Original: I
Revised: I

Original: Initialize belief state z 0 = random numbers, action a 0 = " restart ", and goal g 0 = " <pad> ";
Revised: Initialize belief state z 0 = random numbers, action a 0 = " restart ", and goal g 0 = " <pad> ";

Original: Initialize reward buffer R ; for game id id ∈ I do
Revised: Initialize reward buffer R ; for game id id ∈ I do

Original: Initialize root N 0 =
Revised: Initialize root N 0 =

Original: InitTree ( Q ( z t , a t , g t )) , set t = 0 and R = 0 ; while not end do o t +1 , r t +1 ← env ( a t , id ) and R = R + r t +1 ;
Revised: InitTree ( Q ( z t , a t , g t )) , set t = 0 and R = 0 ; while not end do o t +1 , r t +1 ← env ( a t , id ) and R = R + r t +1 ;

Original: Update object belief z t +1 ∼ q E ( z t +1 | o t +1 , a t , z t ) ;
Revised: Update object belief z t +1 ∼ q E ( z t +1 | o t +1 , a t , z t ) ;

Original: Extract goal: g t +1 =
Revised: Extract goal: g t +1 =

Original: ExtactGoal ( z t +1 , o t +1 , g t ) ;
Revised: ExtactGoal ( z t +1 , o t +1 , g t ) ;

Original: Select action: ˆ a t +1 = MCTS
Revised: Select action: ˆ a t +1 = MCTS

Original: Simulations ( N t , z t +1 , g t +1 , p T ( · ) , p r ( · ) , Q ( · )) ;
Revised: Simulations ( N t , z t +1 , g t +1 , p T ( · ) , p r ( · ) , Q ( · )) ;

Original: Move the search node: N t +1 =
Revised: Move the search node: N t +1 =

Original: Move ( N t , ˆ a t +1 ) ;
Revised: Move ( N t , ˆ a t +1 ) ;

Original: Expand t = t
Revised: Expand t = t

Original: + 1 and end =
Revised: + 1 and end =

Original: CheckEnd ( env ) end
Revised: CheckEnd ( env ) end

Original: Buffer total rewards R = R ∪ R endoutput : Reward buffer R
Revised: Buffer total rewards R = R ∪ R endoutput : Reward buffer R

Original: C C OMPLEMENTARY R ESULTS C. G RAPH V ISUALIZATION
Revised: C C OMPLEMENTARY R ESULTS C. G RAPH V ISUALIZATION

Original: Graph Visualization.
Revised: Graph Visualization.

Original: Figure 7 illustrates examples of the ground-truth memory graph h ∗ t and the graphs learned by the OS-ELBo ( h OSt ) and SS-ELBo ( h SSt ) objectives.
Revised: Figure 7 illustrates examples of the ground-truth memory graph h ∗ t and the graphs learned by the OS-ELBo ( h OSt ) and SS-ELBo ( h SSt ) objectives.

Original: Since the space of adjacency matrices is large ( K × K × C ) while the observed object information is limited, h ∗ t can be represent by a sparse matrix.
Revised: Since the space of adjacency matrices is large ( K × K × C ) while the observed object information is limited, h ∗ t can be represent by a sparse matrix.

Original: Supervised by predicted graphs ˆ h t , OS-ELBo captures the sparsity, and h OSt is a direct approximation to h ∗ t .
Revised: Supervised by predicted graphs ˆ h t , OS-ELBo captures the sparsity, and h OSt is a direct approximation to h ∗ t .

Original: However, graph information are unavailable for the SS-ELBo objective, so h OSt forms a dense latent representation of h ∗ t .
Revised: However, graph information are unavailable for the SS-ELBo objective, so h OSt forms a dense latent representation of h ∗ t .

Original: Similar phenomenon can be observed in
Revised: Similar phenomenon can be observed in

Original: GATA (Adhikari et al., 2020).
Revised: GATA (Adhikari et al., 2020).

Original: We study how well this latent representation and the corresponding object states can represent the object dynamics in Section 4.2.
Revised: We study how well this latent representation and the corresponding object states can represent the object dynamics in Section 4.2.

Original: C. P LANNING E FFICIENCY
Revised: C. P LANNING E FFICIENCY

Original: We compare the efﬁciency of OOTD-based planning algorithms (including Dyna-Q, MCTS and DynaQ+MCTS) by studying which method can gather more rewards with a limited number of selected actions.
Revised: We compare the efﬁciency of OOTD-based planning algorithms (including Dyna-Q, MCTS and DynaQ+MCTS) by studying which method can gather more rewards with a limited number of selected actions.

Original: Figure 8 shows results.
Revised: Figure 8 shows results.

Original: Among the planning algorithms, Dyna-Q+MCTS signiﬁcantly outperforms others.
Revised: Among the planning algorithms, Dyna-Q+MCTS signiﬁcantly outperforms others.

Original: It is because the Q network provides prior knowledge over the preference of actions, based on which MCTS implements Monte-Carlo rollouts to generate and evaluate the belief states of objects.
Revised: It is because the Q network provides prior knowledge over the preference of actions, based on which MCTS implements Monte-Carlo rollouts to generate and evaluate the belief states of objects.

Original: Together, they yield a more accurate look ahead of future rewards.
Revised: Together, they yield a more accurate look ahead of future rewards.