summaries: This paper studies object-oriented text dynamics (OOTD) model and planning in text-based games. OOTD learns internal representation of object dynamics and uses transition layers to predict the belief of object states. Empirical results shows a performance boost compared to model-free baselines. Ablation studies are performed to explain the importance of OOTD components. This work proposes Object-Oriented Text Dynamics for model-based RL in the TextWorld environment. OOTD ensembles existing methods such as ComplEx, R-GCN to build a memory graph, which contains object-level information about each object's attribute and its relations to others, from history observations and predicts states at the next time step. Therefore, the authors can combine the learned model with Dyn-Q or MCTS methods to select actions. If I am correct, under an object-supervised setting, it assumes the existence of a dataset that contains a ground truth memory graph. In contrast, the network has to discover graphs from predicting observation and rewards in the self-supervised setting. Experiments show that OOTD achieves superior performance than various baselines, including model-free RL and previously learned graph models. This paper tackles the difficult problem of learning to play text-based games. This domain is particularly challenging, since the text observations received at each time-step are variable length, and will only provide a partial description of the world, from which the full state must be reconstructed. Additionally, the player must track the history of past states in order to make the correct decision.  This paper introduces a new framework for describing the time-evolution of a game: Object-Oriented Partially Observable Markov Decision Processes (OO POMDP). This work's main contribution is a learned Object-Oriented Text Dynamics (OOTD) model that learns the transition and reward function of a OO POMDP game.  Briefly, the presented method for predicting the transition function is: given a representation $z_t$ of the objects and $a_t$ of the action, create a graph, where the nodes are objects and the edges denote relationships. Perform message passing, and obtain new node representations $e_t$.  Then, (1) given $e_t$, predict a new representation $v_{t}$ using dual stream attention with the action $a_t$. Then, (2) use $v_t$ to predict the updated state $z_{t+1}$ of each object, using an independent set of parameters per object.  The dynamics model is trained using a Evidence lower bound objective (ELBo) in two varieties: one that relies on a deterministic extraction of the graph and one that learns directly from rewards and observations.  To show its effectiveness, the learned dynamics model is used to model-based training (Dyna-Q, MCTS) is used to learn a planner. The learned planner is shown to out-perform the baselines (DQN, DRQN, GATA), both on the test cases, and in terms of sample efficiency. Ablations of (1) and (2) find qualitative advantages of the proposed method, since the learned object representation are both more separable and better semantically clustered. Additionally, the ablations are shown to reduce the accuracy of the graphs and state recreated from the object representations. This paper presents a novel approach for task-based games. The authors used an object-oriented POMDP formulation for the task. The key contribution is to learn the dynamics of the environment with a graph neural network (with either object supervision or self-supervision), and apply an online planning algorithm (e.g., MCTS) to solve the problem. The authors show strong improvements over other approaches in terms of sample efficiency and task scores.
abstractText: The advancement of dynamics models enables model-based planning in complex environments. Dynamics models mostly study image-based games with fully observable states. Generalizing these models to Text-Based Games (TBGs), which often include partially observable states with noisy text observations, is challenging. In this work, we propose an Object-Oriented Text Dynamics (OOTD) model that enables planning algorithms to solve decision-making problems in text domains. OOTD predicts a memory graph that dynamically remembers the history of object observations and filters object-irrelevant information. To improve the robustness of dynamics, our OOTD model identifies the objects influenced by input actions and predicts beliefs of object states with independently parameterized transition layers. We develop variational objectives under the object-supervised and self-supervised settings to model the stochasticity of predicted dynamics. Empirical results show that our OOTD-based planner significantly outperforms model-free baselines in terms of sample efficiency and running scores.
references.authors: Ashutosh Adhikari Xingdi Yuan Marc-Alexandre Côté Mikuláŝ Zelinka Marc-Antoine Rondeau Romain Laroche Pascal Poupart Jian Tang Adam Trischler William L. Hamilton
references.title: Learning dynamic belief graphs to generalize on text-based games
references.venue: In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems
references.authors: Leonard Adolphs Thomas Hofmann
references.title: Ledeepchef deep reinforcement learning agent for families of text-based games
references.venue: In The Thirty-Fourth AAAI Conference on Artificial Intelligence,
references.authors: Prithviraj Ammanabrolu Matthew J. Hausknecht
references.title: Graph constrained reinforcement learning for natural language action spaces
references.venue: In 8th International Conference on Learning Representations,
references.authors: Prithviraj Ammanabrolu Mark Riedl
references.title: Playing text-adventure games with graph-based deep reinforcement learning
references.venue: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,
references.authors: Prithviraj Ammanabrolu Ethan Tien Matthew J. Hausknecht Mark O. Riedl
references.title: How to avoid being eaten by a grue: Structured exploration strategies for textual worlds
references.authors: Christopher P. Burgess Loïc Matthey Nicholas Watters Rishabh Kabra Irina Higgins Matthew Botvinick Alexander Lerchner
references.title: Monet: Unsupervised scene decomposition and representation
references.authors: Marc-Alexandre Côté Akos Kádár Xingdi Yuan Ben Kybartas Tavian Barnes Emery Fine James Moore Matthew Hausknecht Layla El Asri Mahmoud Adada
references.title: Textworld: A learning environment for text-based games
references.venue: In Workshop on Computer Games,
references.authors: Adrien Couëtoux Jean-Baptiste Hoock Nataliya Sokolovska Olivier Teytaud Nicolas Bonnard
references.title: Continuous upper confidence trees
references.venue: Selected Papers,
references.authors: Felipe Leno da Silva Anna Helena Reali Costa
references.title: Object-oriented curriculum generation for reinforcement learning
references.venue: Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems, AAMAS 2018,
references.authors: Felipe Leno da Silva Ruben Glatt Anna Helena Reali Costa
references.title: MOO-MDP: an object-oriented representation for cooperative multiagent reinforcement learning
references.venue: IEEE Trans. Cybern.,
references.authors: Carlos Diuk Andre Cohen Michael L. Littman
references.title: An object-oriented representation for efficient reinforcement learning
references.venue: Machine Learning, Proceedings of the Twenty-Fifth International Conference (ICML
references.authors: Adrien Ecoffet Joost Huizinga Joel Lehman Kenneth O. Stanley Jeff Clune
references.title: Go-explore: a new approach for hard-exploration problems
references.authors: Chelsea Finn Ian J. Goodfellow Sergey Levine
references.title: Unsupervised learning for physical interaction through video prediction
references.venue: Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems
references.authors: Ross B. Girshick
references.title: Fast R-CNN
references.venue: IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13,
references.authors: Vikash Goel Jameson Weng Pascal Poupart
references.title: Unsupervised video object segmentation for deep reinforcement learning
references.venue: Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems
references.authors: Anirudh Goyal Alex Lamb Jordan Hoffmann Shagun Sodhani Sergey Levine Yoshua Bengio Bernhard Schölkopf
references.title: Recurrent independent mechanisms
references.authors: Danijar Hafner Timothy P. Lillicrap Ian Fischer Ruben Villegas David Ha Honglak Lee James Davidson
references.title: Learning latent dynamics for planning from pixels
references.venue: Proceedings of the 36th International Conference on Machine Learning,
references.authors: Matthew J. Hausknecht Ricky Loynd Greg Yang Adith Swaminathan Jason D. Williams
references.title: NAIL: A general interactive fiction agent
references.authors: Matthew J. Hausknecht Prithviraj Ammanabrolu Marc-Alexandre Côté Xingdi Yuan
references.title: Interactive fiction games: A colossal adventure
references.venue: In The Thirty-Fourth AAAI Conference on Artificial Intelligence,
references.authors: Ji He Jianshu Chen Xiaodong He Jianfeng Gao Lihong Li Li Deng Mari Ostendorf
references.title: Deep reinforcement learning with a natural language action space
references.venue: In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics,
references.authors: Vishal Jain William Fedus Hugo Larochelle Doina Precup Marc G. Bellemare
references.title: Algorithmic improvements for deep reinforcement learning applied to interactive fiction
references.venue: In The Thirty-Fourth AAAI Conference on Artificial Intelligence,
references.authors: Mandar Joshi Rakesh Khobragade Saurabh Sarda Umesh Deshpande Shiwali Mohan
references.title: Objectoriented representation and hierarchical reinforcement learning in infinite mario
references.venue: In IEEE 24th International Conference on Tools with Artificial Intelligence,
references.authors: Lukasz Kaiser Mohammad Babaeizadeh Piotr Milos Blazej Osinski Roy H. Campbell Konrad Czechowski Dumitru Erhan Chelsea Finn Piotr Kozakowski Sergey Levine Afroz Mohiuddin Ryan Sepassi George Tucker Henryk Michalewski
references.title: Model based reinforcement learning for atari
references.venue: In 8th International Conference on Learning Representations,
references.authors: Thomas N. Kipf Elise van der Pol Max Welling
references.title: Contrastive learning of structured world models
references.venue: In 8th International Conference on Learning Representations,
references.authors: Levente Kocsis Csaba Szepesvári
references.title: Bandit based monte-carlo planning
references.venue: Machine Learning: ECML 2006, 17th European Conference on Machine Learning, Berlin,
references.authors: Leonid Kuvayev Richard S Sutton
references.title: Model-based reinforcement learning with an approximate, learned model
references.venue: In Proceedings of the ninth Yale workshop on adaptive and learning systems,
references.authors: Nathan O. Lambert Brandon Amos Omry Yadan Roberto Calandra
references.title: Objective mismatch in model-based reinforcement learning
references.venue: Proceedings of the 2nd Annual Conference on Learning for Dynamics and Control,
references.authors: Andrea Madotto Mahdi Namazifar Joost Huizinga Piero Molino Adrien Ecoffet Huaixiu Zheng Alexandros Papangelis Dian Yu Chandra Khatri Gökhan Tür
references.title: Exploration based language learning for text-based games
references.venue: In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence,
references.authors: Ofir Marom Benjamin Rosman
references.title: Zero-shot transfer with deictic object-oriented representation in reinforcement learning
references.venue: Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems
references.authors: Shiwali Mohan John E. Laird
references.title: An object-oriented approach to reinforcement learning in an action game
references.venue: Proceedings of the Seventh AAAI Conference on Artificial Intelligence and Interactive Digital Entertainment, AIIDE 2011, October
references.authors: Karthik Narasimhan Tejas D. Kulkarni Regina Barzilay
references.title: Language understanding for textbased games using deep reinforcement learning
references.venue: Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing,
references.authors: Tom Schaul Daniel Horgan Karol Gregor David Silver
references.title: Universal value function approximators
references.venue: Proceedings of the 32nd International Conference on Machine Learning, ICML 2015,
references.authors: Michael Sejr Schlichtkrull Thomas N. Kipf Peter Bloem Rianne van den Berg Ivan Titov Max Welling
references.title: Modeling relational data with graph convolutional networks
references.venue: The Semantic Web - 15th International Conference,
references.authors: Min Joon Seo Aniruddha Kembhavi Ali Farhadi Hannaneh Hajishirzi
references.title: Bidirectional attention flow for machine comprehension
references.venue: In 5th International Conference on Learning Representations,
references.authors: David Silver Aja Huang Chris J. Maddison Arthur Guez Laurent Sifre George van den Driessche Julian Schrittwieser Ioannis Antonoglou Vedavyas Panneershelvam Marc Lanctot Sander Dieleman Dominik Grewe John Nham Nal Kalchbrenner Ilya Sutskever Timothy P. Lillicrap Madeleine Leach Koray Kavukcuoglu Thore Graepel Demis Hassabis
references.title: Mastering the game of go with deep neural networks and tree
references.venue: search. Nat.,
references.authors: David Silver Thomas Hubert Julian Schrittwieser Ioannis Antonoglou Matthew Lai Arthur Guez Marc Lanctot Laurent Sifre Dharshan Kumaran Thore Graepel
references.title: A general reinforcement learning algorithm that masters chess, shogi, and go through self-play
references.authors: Ilya Sutskever Oriol Vinyals Quoc V. Le
references.title: Sequence to sequence learning with neural networks
references.venue: Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems
references.authors: Richard S Sutton Andrew G Barto
references.title: Reinforcement learning: An introduction
references.venue: MIT press,
references.authors: Naftali Tishby Noga Zaslavsky
references.title: Deep learning and the information bottleneck principle
references.venue: IEEE Information Theory Workshop,
references.authors: Adam Trischler Marc-Alexandre Côté Pedro Lima
references.title: First TextWorld Problems, the competition: Using text-based games to advance capabilities of AI agents
references.authors: Théo Trouillon Johannes Welbl Sebastian Riedel Éric Gaussier Guillaume Bouchard
references.title: Complex embeddings for simple link prediction
references.venue: Proceedings of the 33nd International Conference on Machine Learning,
references.authors: Jack Urbanek Angela Fan Siddharth Karamcheti Saachi Jain Samuel Humeau Emily Dinan Tim Rocktäschel Douwe Kiela Arthur Szlam Jason Weston
references.title: Learning to speak and act in a fantasy text adventure game
references.venue: In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing,
references.authors: Laurens Van der Maaten Geoffrey Hinton
references.title: Visualizing data using t-SNE
references.venue: Journal of machine learning research,
references.authors: Arthur Wandzel Yoonseon Oh Michael Fishman Nishanth Kumar Lawson L.S. Wong Stefanie Tellex
references.title: Multi-object search using object-oriented pomdps
references.venue: In International Conference on Robotics and Automation,
references.authors: Tingwu Wang Xuchan Bao Ignasi Clavera Jerrick Hoang Yeming Wen Eric Langlois Shunshi Zhang Guodong Zhang Pieter Abbeel Jimmy Ba
references.title: Benchmarking model-based reinforcement learning
references.authors: Xusen Yin Jonathan May
references.title: Comprehensible context-driven text game playing
references.venue: In IEEE Conference on Games, CoG 2019,
references.authors: Xingdi Yuan Marc-Alexandre Côté Alessandro Sordoni Romain Laroche Remi Tachet des Combes Matthew J. Hausknecht Adam Trischler
references.title: Counting to explore and generalize in text-based games
references.authors: Tom Zahavy Matan Haroush Nadav Merlis Daniel J. Mankowitz Shie Mannor
references.title: Learn what not to learn: Action elimination with deep reinforcement learning
references.venue: Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems
references.authors: Mikuláŝ Zelinka
references.title: Baselines for reinforcement learning in text games
references.venue: IEEE 30th International Conference on Tools with Artificial Intelligence, ICTAI 2018,
references.authors: Guangxiang Zhu Zhiao Huang Chongjie Zhang
references.title: Object-oriented dynamics predictor
references.venue: Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems
references.authors: Guangxiang Zhu Jianhao Wang Zhizhou Ren Zichuan Lin Chongjie Zhang
references.title: Object-oriented dynamics learning through multi-level abstraction
references.venue: In The Thirty-Fourth AAAI Conference on Artificial Intelligence,
references.authors: Adhikari
references.title: 2020), which show that the agent can properly solve a game by modelling these objects and relations. The learning rate of policy training and dynamics model training is set to 0.0001. The discount factor γ is set to 0.9. The size of the hidden layers and the size of z in our OOTD model are set to 32. The transition layers of different objects are parameterized independently, so the neural network predicts 32×
references.authors: Adhikari
references.title: Under review as a conference paper at ICLR 2022 games in a cooking theme across a wide range of difficulty levels. For each game, the expert trajectory and the memory graph ht at each time step are recorded. The original dataset has only 10 games for each difficulty level. This is insufficient for dynamics model pre-training
sections.heading: 1 INTRODUCTION
sections.text: Planning algorithms typically leverage the environment dynamics to solve decision-making problems (Sutton & Barto, 2018). To plan in unknown environments, the agent must learn a dynamics model to predict future (belief) states and rewards by conditioning on an action. This dynamics model enables the implementation of intensive search for optimal actions, which can potentially increase both sample efficiency and cumulative rewards, compared to model-free methods (Hafner et al., 2019; Wang et al., 2019; Kaiser et al., 2020). 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.
  To facilitate dynamics learning in complex environments, Diuk et al. (2008) proposed an ObjectOriented Markov Decision Process (OO-MDP) that factorizes world states into object states. They showed that the agent can find optimal policies with a better sample efficiency by modeling the dynamics at the object level. Some following works (Finn et al., 2016; Goel et al., 2018; Zhu et al., 2018; 2020) extended OO-MDPs to image-based games. These methods typically assume full observability over game states and a fixed input size, which facilitates the use of object masks to decompose an image into different objects. 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.51% of the candidate objects are mentioned in an observation from Textworld (Côté et al., 2018)). The dynamics model must remember the history of observed objects to predict accurately their states. Moreover, an observation typically contains lots of noisy patterns that record object-irrelevant information (e.g., the non-bolded text from ot in Figure 1). 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.
  In this work, we design an Object-Oriented Text Dynamics (OOTD ) model that integrates:
  1) Graph Representation for Objects. OOTD predicts a memory graph from the input states. The memory graph, whose nodes correspond to the objects to be modelled, captures the information about these objects from the beginning of a game. By applying this graph as an information bottleneck for
  dynamics prediction, OOTD filters the object-irrelevant information in the text inputs, which forms a tighter information plane and facilitates generalization to unseen games (Tishby & Zaslavsky, 2015).
  2) Independent Transition Layers. Intuitively, an input action influences 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. To capture which objects are influenced 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. This independent mechanism enhances the robustness for dynamics prediction (Goyal et al., 2019) (also see the ablation study in Section 4.2).
  During planning, we predict belief states without knowing rewards and observations. To learn a stochastic dynamics model, we propose object-supervised and self-supervised Evidence Lower Bound (ELBo) objectives for training our OOTD model. We evaluate how well the OOTD model supports planning by implementing planning algorithms (e.g., Dyna-Q (Kuvayev & Sutton, 1996) and Monte Carlo Tree Search (MCTS) (Kocsis & Szepesvári, 2006)) based on learned dynamics. Empirical results show that these planning algorithms significantly outperform other baselines in terms of game performance and sample efficiency. To support the design of our OOTD model, we include an ablation study that visualizes the object states and quantifies the dynamics prediction performance.
  Contributions: 1) We introduce our approach to implementing model-based planning algorithms to solve decision-making problems in the text-domain. 2) We propose the OOTD model that learns to predict beliefs about object dynamics.
sections.heading: 2 OBJECT LEVEL PLANNING IN TEXT DOMAINS
sections.text: 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. We define a graph representation to remember object information and introduce the corresponding planning algorithms.
sections.heading: 2.1 TEXT-BASED GAMES
sections.text: 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. 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. The distributions of rewards over state-action pairs are different in training and testing games because TBGs assign rewards by conditioning on specific goals (that are predefined, but unknown). 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. 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. This property enables the generalization of knowledge learned from training games to solve testing games.
sections.heading: 2.2 OBJECT-ORIENTED PARTIALLY OBSERVABLE MARKOV DECISION PROCESSES
sections.text: Since an observation reveals only partial information about the game state, we formulate TBGs as partially observable Markov decision processes (POMDPs). In this work, we extend the POMDP framework to represent object-level dynamics (Diuk et al., 2008). The Object-Oriented (OO) POMDP is a tuple 〈S,O,Z,Φ,G,A,R, T , γ〉, where:
  • S andO define the spaces of low-level states and observations from the TGB environments while Z and Φ are the spaces of object-level states and observations. In a TBG, o ∈ O is a text message from the game environment (e.g., "Please fry the potato."), whereas φ ∈ Φ records the specific objects and their relations in this message (e.g., a triplet like "potato-need-fry")). To model the object dynamics, the agent must extract φ from o by distilling objects’ information from text sentences. In this work, we model object states z ∈ Z (instead of s ∈ S) and learn latent representations for z.
  • G and A are the spaces for goals and actions. We study choice-based games (Yin & May, 2019), where candidate actions At ∈ A are available at time step t. We include a goal variable g to mark different tasks in each game. 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 finished.
  • R and T define the spaces of reward and transition models. γ is a discount factor. 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, . . . ,K}, pkT (zk,t|zt−1, at−1) ∈ T and pr(rt|zt−1, gt) ∈ R.
  In our OO-POMDP, transitions pT , observations φ, and states z are modelled for each object, whereas actions a and rewards r are defined for the entire environment. This generalizes to popular RL environments that accept an action and return a reward at every time step. 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).
sections.heading: 2.3 GRAPH REPRESENTATION FOR OBJECTS
sections.text: We utilize an object-relation-object triplet φ to represent object information. At time step t, φt = [φt,1, ..., φt,M ] (M is the number of observed triplets at t). Given a total of K candidate objects and C candidate relations, these triplets can be mapped to a knowledge graph Ωt ∈ {0, 1}C×K×K . Each entry (c, i, j) ∈ {0, 1} indicates whether there is a relationship r between the ith and jth objects. This knowledge graph forms a natural representation of object relations since object information in most observations ot corresponds to either entity attributes (e.g., "potato-is-fried") or to relational information about entities(e.g., "potato-on-counter" and "bedroom-north_of-kitchen"), which can be conveniently represented by triplets and the corresponding graph.
  Memory Graph. We store the observations up to time step t in a memory graph ht = Ω0 ⊕ Ω1 ⊕ · · · ⊕ Ωt that captures object information observed by the agent since the beginning of a game. Similar to previous works Adhikari et al. (2020); Ammanabrolu & Riedl (2019); Ammanabrolu & Hausknecht (2020), we summarize the object-level history with a latent memory graph ht. To update the memory from t − 1 to t, we learn a graph updater ⊕, and ht = ht−1 ⊕ Ωt. 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). Our OOTD model is trained to automatically emulate such an operator. Our transition model (Section 3.1) is trained to generate ht from zt, allowing the latent object states to capture object relations.
sections.heading: 2.4 MODEL-BASED PLANNING IN TEXT-BASED GAMES
sections.text: Table 1: Object dynamics for planning, where the transition models and observation encoders are independently parameterized for a total of K objects.
  Transition models
  [
  pkT (zk,t+1|at, zt)
  ]K
  k=1
  Observation encoders [ qkE(zk,t+1|ot+1, at, zt) ]K k=1 Reward model pr(rt|zt, gt) Graph & Obs. decoder pΩ(ht|zt), po(ot|zt)
  We introduce model-based planning in TBGs. Based on the OO-POMDP (Section 2.2), at each time step t, we define latent states zt = [z1,t, . . . , zK,t] for a total of K objects, text observations ot, action commands at, and scalar rewards rt, that follow the stochastic dynamics in Table 1. The details of these dynamics models are introduced in Section 3.
  Based on these dynamics, we implement planning algorithms to select an action from candidate commands At for maximizing the expected sum of rewards Eπ( ∑T t=0 γ
  trt), as shown in Figure 1. In this work, we study choice-based games (Yin & May, 2019), where the candidate commands (or actions) At are available and the planner determines the action at ∈ At to be performed. The planning algorithms include Dyna-Q, Monte-Carlo Tree Search (MCTS) and their combinations:
  Dyna-Q (Kuvayev & Sutton, 1996) incorporates dynamics models and Q-learning. 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. Compared to the model-free Deep Q-Network (DQN), Dyna-Q is more sample efficient: by expanding the replay buffer with dynamics models, Dyna-Q converges faster with the same number of interactions from the environment.
  LTD = E [( rt + γ max
  a∈At Q(zt, a, gt)−Q(zt−1, at−1, gt−1)
  )2] (1)
  MCTS (Kocsis & Szepesvári, 2006) is a heuristic search algorithm that builds and updates a search tree based on environment dynamics. By performing Monte-Carlo simulations organized in a tree structure, MCTS does an efficient search in environments with large action spaces (Silver et al., 2018). MCTS iteratively runs multiple playouts. At the ith 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 ai,t to maximize the Upper Confidence Bound (UCB) (Couëtoux et al., 2011):
  ai,t = arg max a∈At
  [ Qi(zt, a, g) + cpuct √ log(i− 1)
  ζi−1(zt, a, g) + 1
  ] (2)
  where cpuct controls the scale of exploration, and ζi records the visit count at the ith play. 2) Evaluation: Evaluate the selected leaf zτ with the reward model pr(rτ |zτ , gt). 3) Expansion: Expand the leaf node by adding child nodes. 4) Back Up: Update the action-values: Qi = (Qi−1ζi−1 + rτ )/(ζi−1 + 1) and increment the visit count: ζi = ζi−1 + 1 on all the traversed edges.
  Dyna-Q + MCTS initializes Q values in MCTS by utilizing the Dyna-Q network. This combination enables MCTS to perform the tree search with more efficiency and efficacy (Silver et al., 2016).
sections.heading: 3 OBJECT-ORIENTED DYNAMICS MODELS
sections.text: 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. 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. We introduce the object-supervised and self-supervised objectives to train these models.
sections.heading: 3.1 OBJECT ORIENTED TRANSITION MODEL
sections.text: Figure 2 illustrates our transition model, which has three major components (a) a graph decoder, (b) a graph encoder, and (c) independent transition layers.
  (a) Graph Decoder. Given the states of K objects zt−1 = [z1,t−1, ..., zK,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 ht−1 ∈ [0, 1]C×K×K . The scoring function for an entry (c, i, j) (in the adjacency matrix of ht−1) is implemented by Re(〈ωc, zi,t−1, zj,t−1〉) where ωc ∈ CE is a complex vector for the relation c and Re denotes the real part of the decomposition in complex space. The function
  can model a large variety of binary relations, including both symmetric and antisymmetric relations. Based on this scoring function, predicting ht−1 is equivalent to estimating the probability of having a relation c between object i and j, which can be efficiently approximated with a low-ranked (E K) matrix decomposition: ht = [Re(ZW1ZT ), ..., Re(ZWCZT )] where Z ∈ RK×E is the matrix of object states andWc ∈ CE×E is a complex matrix. In practice, this decomposition enables to map efficiently from zt−1 to ht−1, therefore scaling our method to environments with numerous objects.
  (b) Graph Encoder. We encode the memory graph ht−1 into node representations et−1 = [e1,t−1, ..., eK,t−1] with a Relational Graph Convolution Network (R-GCN) (Schlichtkrull et al., 2018). Nodes in ht correspond to candidate objects and RGCN performs message passing between nodes, so ek,t−1 captures the semantic information of the kth object and its relations to other nodes (objects). Compared to zt−1, et−1 denotes the object embeddings after taking into account relations to other objects in the graph. The graph can be thought as a succinct latent representation that ignores irrelevant object information in text observations when predicting the next state zt.
  (c) Independent Transition Layers. TBG environments typically contain a large population of objects, and each action can only affect a few underlying objects. 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).
  Given et−1 and an action at−1 = [a1,t−1, ..., aJ,t−1] (a sentence of J words), the BIDAF network learns an attention matrixBae ∈ [0, 1]K×J , where bk,j measures the similarity between the representation of the k-th object and the embedding of the j-th word. To learn an action-aware representation of object, we compute 1) action-to-object attentions be 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 bak that indicate which action words are most relevant to the kth object representation. The action-aware object representation concatenates the attended action vector νak,t−1, the attended object representation ν e t−1 and the initial object representation
  ek,t−1: νak,t−1 = ∑
  j bak,jψ a(aj,t−1) where bak = softmax(B ae k,:) ∈ [0, 1]J , νet−1 = ∑
  k bekek,t−1 where b e = softmax(maxcol(B ae)) ∈ [0, 1]K ,
  pkT (zk,t|at−1, zt−1) = N (µk,t,σk,t) where [µk,t,σk,t] = ψ p k([ν a k,t−1,ν e t−1, ek,t−1]) (3) Here ψa is a transformer-based text embedding function (Vaswani et al., 2017) for the input action, and ψpk is implemented by a MLP. The transition layers ψ p 1 , . . . , ψ p K are parameterized independently.
sections.heading: 3.2 OBJECT ORIENTED REWARD MODEL
sections.text: We learn a reward model pr(rt|zt, gt) based on the predicted belief states. 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.
  To implement the reward model, we build a goal extractor (Appendix A.3) that extracts goals gt from object states zt = [z1,t, . . . , zK,t]. 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. Although zt implicitly embeds gt with latent values, we prefer to represent gt explicitly with a sentence for the ease of model comprehension. Given the goal gt and object states zt, the reward model pr is defined by
  pr(rt|zt, gt) = ψr[zpoolt , ψg(gt)] where z pool t = mean_pooling[z1,t, . . . , zK,t], (4)
  ψg is a transformer-based text embedding function for goals, and ψr is implemented by a MLP.
sections.heading: 3.3 TRAINING OBJECTIVE FOR DYNAMICS MODELS
sections.text: 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.
  Object-Supervised Setting. We learn the dynamics model under the supervision of memory graphs ht that record rich information about the observed objects and their relations from the beginning of a game (Section 2.3). The Object-Supervised Evidence Lower Bound (OS-ELBo) objective is:
  T∑ t=1 { EqE [ log pΩ(ĥt|zt) ] − K∑ k=1 Dkl [ qkE(zk,t|ot, at−1, zt−1)‖pkT (zk,t|at−1, zt−1) ]} (5)
  where pΩ is the graph decoder, pT is the transition model and qkE is the observation encoder. We now provide more details about each component and how to obtain ht. pΩ is approximated by a re-sampling technique described in Appendix A.2. pT is a Gaussian distribution as described in Equation 3. qkE shares a similar structure to pT (see Appendix A.5 for more details).
  The memory graph ht is not a direct output from the RL environment. To perform object-supervised training, we extract object information from ot and predict ĥt with a deterministic object extractor fe(ot, at−1, ht−1). We learn fe during a pre-training stage by utilizing a FTWP dataset. The dataset records the trajectories [o1, a1, . . . , ot, at] and corresponding [h1, . . . , ht] by following walkthroughs in FTWP games. In practice, we can build such a dataset by extracting triplets (φ) from the text observations (o1, . . . , ot) with an Open Information Extractor (OIE) (Angeli et al., 2015) and map the captured triplets into a knowledge graph by following (Ammanabrolu & Hausknecht, 2020). After training fe, we initialize h0 = {0}C×K×K and update by ĥt = fe(ot, at−1, ĥt−1) during training. The detailed introduction of fe and the FTWP dataset can be found in Appendix A.4 and A.6.
  To train the reward model, we sample object states zt = [z1,t, . . . , zK,t] from their belief states, extract goal gt, and minimize a smooth L1 loss (Girshick, 2015):{
  1 2β [pr(rt|zt, gt)− rt] 2, if |pr(rt|zt, gt)− rt| < β, |pr(rt|zt, gt)− rt| − 12β, otherwise.
  (6)
  Self-Supervised Setting. 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 Evidence Lower Bound (SS-ELBo) objective that directly learns from rewards and observations: (7)T∑ t=1 { EqE [ log po(ot|zt) + log pr(rt|zt, gt) ] − K∑ k=1 Dkl [ qE(zk,t|ot, at−1, zt−1)‖pT (zk,t|at−1, zt−1)
  ]} where pT , pr and qE are the aforementioned transition model, reward model and observation encoder. The observation decoder po is implemented by a Sequence-to-Sequence (Seq2Seq) model (Sutskever et al., 2014), which is trained by a teacher-forcing technique with text observations ot. This SS-EBLo objective enables training the reward and transition model together by maximizing one objective, which improves the training efficiency. Appendix C.1 shows an example of a predicted graph.
sections.heading: 4 EXPERIMENTS
sections.text: Environments. We divide the games in the Text-World benchmark (Côté et al., 2018) into five subsets according to their difficulty levels. Each subset contains 100 training, 20 validation, and 20 testing games. These subsets are mutually exclusive. 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. Table 4 summarizes the game statistics.
  Running Settings. To efficiently 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 efficiency with other model-free baselines. After training, we fix the parameters in dynamics models and predict belief states and rewards to solve the games in the test dataset. 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). The model hyper-parameters are tuned with the games in the validation set. (Appendix A.1).
  Baselines. we compare the state-of-the-art models for solving TBGs. The baselines include Deep Q-Network (DQN) (Narasimhan et al., 2015) and Deep Recurrent Q-Network (DRQN) (Yuan et al., 2018). DQN selects the action based on the current observation, whereas DRQN models the history of action and observations with an RNN. For fair comparisons, we replace their LSTM-based text encoders with transformers. We compare an extension of DRQN by using an episodic counting bonus
  to encourage exploration, which is denoted as (DRQN+) (Yuan et al., 2018). Our next baseline is KGA2C (Ammanabrolu & Hausknecht, 2020), which builds a graph-based state space and applies the Actor-Critic algorithm for learning the policy. We replace their action generator with an action selector for comparison. 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. 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). Note that we apply the same pre-training dataset for our OS-ELBo objective. 2) GATA-OG that learns a continuous graph updater by reconstructing the text observation. 3) GATA-COC that trains the graph updater with a contrastive observation classification objective. Both GATA-OG and GATA-COC follow the self-supervised setting, which offers fair comparison to our SS-ELBo objective.
sections.heading: 4.1 CONTROLLING PERFORMANCE IN TBG ENVIRONMENTS
sections.text: Table 2 shows the agent’s controlling performance at different difficulty levels. Equipped with the OOTD model, our agent implements model-based planning and achieves leading performance over the agents based on model-free algorithms. We find the SS-based OOTD models outperform the OS-based models by scoring more rewards at 4 (out of 6) difficulty levels. Although the object signals facilitate a better prediction of dynamics, this advantage can not be directly transferred to controlling performance. This is commonly known as the objective-mismatch problem in modelbased RL (Lambert et al., 2020). On the other hand, SS-based models are directly optimized toward generating a better prediction of rewards, which are important signals for controlling. The rewardirrelevant information is filtered to derive a more compressed representation for better generalization ability. Another important finding is that our OOTD-based agents perform better than other graphbased agents (e.g., GATA-GTP and KG-A2C). 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.
  Sample Efficiency. We empirically demonstrate the sample efficiency of our model-based agents by comparing their training performance with other model-free agents. 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. 3) Object-Oriented (OO)-DQN that trains the OOTD model to predict object states, based on which we compute action values. As a model-free method, OO-DQN does not expand the replay buffer with the samples generated by dynamics models. Figure 3 shows the training curve for 6 difficulty levels. 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 difficulty level 5, where GATA-OG converges faster. We find it generally takes more samples to learn the interactions among K objects and their relations to action values. It explains
  why (MF)-Dyna-Q converges slower than others. However, the object-based value function performs better than that based on raw observation in terms of normalized scores.
sections.heading: 4.2 ABLATION STUDIES FOR OBJECT-ORIENTED DYNAMICS
sections.text: 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., 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)). We compare these models with our OOTD models trained by the OS-ELBo objective (OS-OOTD) and the SSELBo objective (SS-OOTD). To evaluate how well the modified model pkT ′ captures object dynamics, we sample object states zk,t ∼ pkT ′(zk,t|at−1, zt−1) for a total of K objects, visualize these states with the t-distributed Stochastic Neighbor Embedding (t-SNE) algorithm (Van der Maaten & Hinton, 2008), and quantify the dynamics prediction performance.
  Object States Visualization. In this experiment, we randomly pick 20 games from the testing games of all difficulty levels. Within each game, pkT ′ computes the states zk,1, · · · , zk,T with randomly selected actions a0, . . . , aT−1, and t-SNE embeds zk,t into a vector of 2 dimensions. Figure 4 visualizes the embedded vectors labelled by 10 randomly selected objects. 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. This phenomenon is not observable when we apply a single layer for computing object states (i.e., SLa-Dyna). Another important observation is that the OOTD model automatically learns to generate similar states for similar objects. For example, OS-OOTD generates similar states for the object of ingredients (e.g., parsley, apple, carrot, and potato marked by the dark frame). A similar phenomenon, although less obvious, can be observed for SS-OOTD. However, such a shrinkage effect disappears in NAt-Dyna, which explains the importance of bidirectional attentions.
  Probing Experiment. This experiment studies how well our OOTD models can capture useful information from the TBG environments. We map the sampled object states z1,t, . . . , zK,t to memory graphs ĥt and rewards r̂t with the additional graph decoder p′G and reward model p ′ r. The performance
  of the dynamics models is quantified by whether ĥt and r̂t can match the ground-truth h∗t and r ∗ t in the FTWP dataset (Appendix A.6). During sampling, the input actions a1, . . . , aT follow the recorded trajectories in FTWP. The ground-truth rewards in FTWP games is either one or zero, and we perform an entry-level evaluation (e.g., ht[i, j, c]) on the adjacency matrix of a predicted graph. It allows us to evaluate whether the real rewards and triplets are captured with F1 scores. 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 first freeze the parameters in the dynamics models and then update the added p′G and p ′ r with data in training games.
  Table 3 shows the testing performance. 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. We add a Real and a Random baseline that feed the groundtruth graph h∗t and a randomly generated graph h RM t into the graph encoder of transition models. They study the probing performance when ground-truth graphs are known and when the posterior collapses for the transition model (so pkT ′(zk,t|at−1, zt−1) = pkT ′(zk,t)). The results show that OS-OOTD achieves a better graph prediction performance by applying the object-supervised objectives (OS-ELBo), but SS-
  OOTD has higher reward prediction accuracy, since SS-ELBo includes an end-to-end training to predict future rewards, which filters the reward-irrelevant object information. Another observation is that removing independent layers and directional attentions significantly degrade the graph prediction performance. It demonstrates that they are important components of our OOTD model.
sections.heading: 5 RELATED WORK
sections.text: DRL Agent for Text-Based Games (TBGs). TBGs (e.g., Textworld (Côté et al., 2018) and interactive fictional games (Hausknecht et al., 2020)) are interactive simulations with text observations and actions. 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), 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). Since states in the TBGs environment are partially observable, many previous studies have formulated TBGs as POMDPs and modelled the history of observations with a Deep Recurrent Q-Network (DRQN) (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). Previous works are commonly based on model-free Q learning while we explore model-based planning.
  Object Oriented Reinforcement Learning. Diuk et al. (2008) proposed modelling the transition and reward models at the object level. 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. However, object features are not commonly available from the RL environments. Recent works have extracted object representations from states with optical flow (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). Previous works commonly studied image-based fully observable environments whereas we explore text-based partially observable environments.
sections.heading: 6 CONCLUSION
sections.text: We proposed an OOTD model that enables planning in text-based environments. The OOTD model learned a graph representation for capturing object dynamics and predicted the belief of object states with independent transition layers. Empirical results demonstrate that OOTD-based agents outperform model-free baselines in terms of controlling performance and sample efficiency. We then provided an ablation study to show the importance of OOTD components. 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.
sections.heading: 7 ETHICS STATEMENT
sections.text: In this work, we study model-based planning in a simulated (TBG) environment. Although the task itself has limited consequences for society, we take a broader view of our work and discuss its potential influence on future research that might have a considerable impact on our society. 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 specific tasks. 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. During this process, failing to communicate clearly, sensibly, or convincingly might also cause harm, for example, without sufficient calibration, the machine could generate improper words or phrases that might potentially cause a bias toward some minorities or underrepresented groups. How to eliminate such bias will be an important direction of future research.
sections.heading: 8 REPRODUCIBILITY STATEMENT
sections.text: This section introduces how we facilitate the reproducibility from the following perspectives:
  Model Implementation. Section 4 introduces our approach to dividing the training, validation, and testing datasets as well as the running settings based on the constructed datasets. We report the hyper-parameters for training our OOTD model and the seeds for independent runs in Appendix A.1. The computing resources and running time are discussed in Appendix A.7.
  Code and Dataset. We uploaded the code to the submission system. A public version will be added to GitHub once the anonymous period ends. We include all the required packages and environments settings to implement our model in the text document of the uploaded code. In terms of the dataset, we used public and self-generated datasets. For the public one, we provide an introduction in Appendix A.6 and the corresponding download link. For the self-generated one, we have uploaded the dataset to our net drive. The download link will be included in the uploaded code once the anonymous period ends.
sections.heading: Association for Computational Linguistics and the 7th International Joint Conference on Natural
sections.heading: Language Processing of the Asian Federation of Natural Language Processing, ACL 2015, July
sections.text: 26-31, 2015, Beijing, China, Volume 1: Long Papers, pp. 344–354. The Association for Computer Linguistics, 2015.
  Christopher P. Burgess, Loïc Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matthew Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. CoRR, abs/1901.11390, 2019.
  Marc-Alexandre Côté, Akos Kádár, Xingdi Yuan, Ben Kybartas, Tavian Barnes, Emery Fine, James Moore, Matthew Hausknecht, Layla El Asri, Mahmoud Adada, et al. Textworld: A learning environment for text-based games. In Workshop on Computer Games, pp. 41–75. Springer, 2018.
  Adrien Couëtoux, Jean-Baptiste Hoock, Nataliya Sokolovska, Olivier Teytaud, and Nicolas Bonnard. Continuous upper confidence trees. In Carlos A. Coello Coello (ed.), Learning and Intelligent Optimization - 5th International Conference, LION 5, Rome, Italy, January 17-21, 2011. Selected Papers, volume 6683 of Lecture Notes in Computer Science, pp. 433–445. Springer, 2011.
  Felipe Leno da Silva and Anna Helena Reali Costa. Object-oriented curriculum generation for reinforcement learning. In Elisabeth André, Sven Koenig, Mehdi Dastani, and Gita Sukthankar (eds.), Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems, AAMAS 2018, Stockholm, Sweden, July 10-15, 2018, pp. 1026–1034. International Foundation for Autonomous Agents and Multiagent Systems Richland, SC, USA / ACM, 2018.
  Felipe Leno da Silva, Ruben Glatt, and Anna Helena Reali Costa. MOO-MDP: an object-oriented representation for cooperative multiagent reinforcement learning. IEEE Trans. Cybern., 49(2): 567–579, 2019.
  Carlos Diuk, Andre Cohen, and Michael L. Littman. An object-oriented representation for efficient reinforcement learning. In William W. Cohen, Andrew McCallum, and Sam T. Roweis (eds.), Machine Learning, Proceedings of the Twenty-Fifth International Conference (ICML 2008), Helsinki, Finland, June 5-9, 2008, volume 307 of ACM International Conference Proceeding Series, pp. 240–247. ACM, 2008.
  Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O. Stanley, and Jeff Clune. Go-explore: a new approach for hard-exploration problems. CoRR, abs/1901.10995, 2019.
  Chelsea Finn, Ian J. Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 64–72, 2016.
  Ross B. Girshick. Fast R-CNN. In 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 1440–1448. IEEE Computer Society, 2015.
  Vikash Goel, Jameson Weng, and Pascal Poupart. Unsupervised video object segmentation for deep reinforcement learning. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 5688–5699, 2018.
  Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. CoRR, abs/1909.10893, 2019.
  Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2555–2565. PMLR, 2019.
  Matthew J. Hausknecht, Ricky Loynd, Greg Yang, Adith Swaminathan, and Jason D. Williams. NAIL: A general interactive fiction agent. CoRR, abs/1902.04259, 2019.
  Matthew J. Hausknecht, Prithviraj Ammanabrolu, Marc-Alexandre Côté, and Xingdi Yuan. Interactive fiction games: A colossal adventure. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 7903–7910. AAAI Press, 2020.
  Ji He, Jianshu Chen, Xiaodong He, Jianfeng Gao, Lihong Li, Li Deng, and Mari Ostendorf. Deep reinforcement learning with a natural language action space. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, ACL 2016, August 7-12, 2016, Berlin, Germany, Volume 1: Long Papers. The Association for Computer Linguistics, 2016.
  Vishal Jain, William Fedus, Hugo Larochelle, Doina Precup, and Marc G. Bellemare. Algorithmic improvements for deep reinforcement learning applied to interactive fiction. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 4328–4336. AAAI Press, 2020.
  Mandar Joshi, Rakesh Khobragade, Saurabh Sarda, Umesh Deshpande, and Shiwali Mohan. Objectoriented representation and hierarchical reinforcement learning in infinite mario. In IEEE 24th International Conference on Tools with Artificial Intelligence, ICTAI 2012, Athens, Greece, November 7-9, 2012, pp. 1076–1081. IEEE Computer Society, 2012.
  Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
  Thomas N. Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
  Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In Johannes Fürnkranz, Tobias Scheffer, and Myra Spiliopoulou (eds.), Machine Learning: ECML 2006, 17th European Conference on Machine Learning, Berlin, Germany, September 18-22, 2006, Proceedings, volume 4212 of Lecture Notes in Computer Science, pp. 282–293. Springer, 2006.
  Leonid Kuvayev and Richard S Sutton. Model-based reinforcement learning with an approximate, learned model. In Proceedings of the ninth Yale workshop on adaptive and learning systems, pp. 101–105, 1996.
  Nathan O. Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. In Alexandre M. Bayen, Ali Jadbabaie, George J. Pappas, Pablo A. Parrilo, Benjamin Recht, Claire J. Tomlin, and Melanie N. Zeilinger (eds.), Proceedings of the 2nd Annual Conference on Learning for Dynamics and Control, L4DC 2020, Online Event, Berkeley, CA, USA, 11-12 June 2020, volume 120 of Proceedings of Machine Learning Research, pp. 761–770. PMLR, 2020.
  Andrea Madotto, Mahdi Namazifar, Joost Huizinga, Piero Molino, Adrien Ecoffet, Huaixiu Zheng, Alexandros Papangelis, Dian Yu, Chandra Khatri, and Gökhan Tür. Exploration based language learning for text-based games. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI 2020, pp. 1488–1494. ijcai.org, 2020.
  Ofir Marom and Benjamin Rosman. Zero-shot transfer with deictic object-oriented representation in reinforcement learning. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 2297–2305, 2018.
  Shiwali Mohan and John E. Laird. An object-oriented approach to reinforcement learning in an action game. In Vadim Bulitko and Mark O. Riedl (eds.), Proceedings of the Seventh AAAI Conference on Artificial Intelligence and Interactive Digital Entertainment, AIIDE 2011, October 10-14, 2011, Stanford, California, USA. The AAAI Press, 2011.
  Karthik Narasimhan, Tejas D. Kulkarni, and Regina Barzilay. Language understanding for textbased games using deep reinforcement learning. In Lluís Màrquez, Chris Callison-Burch, Jian Su, Daniele Pighin, and Yuval Marton (eds.), Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, EMNLP 2015, Lisbon, Portugal, September 17-21, 2015, pp. 1–11. The Association for Computational Linguistics, 2015.
  Judea Pearl. Causality. Cambridge university press, 2009.
  Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, 2015.
  Michael Sejr Schlichtkrull, Thomas N. Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. Modeling relational data with graph convolutional networks. In Aldo Gangemi, Roberto Navigli, Maria-Esther Vidal, Pascal Hitzler, Raphaël Troncy, Laura Hollink, Anna Tordai, and Mehwish Alam (eds.), The Semantic Web - 15th International Conference, ESWC 2018, Heraklion, Crete, Greece, June 3-7, 2018, Proceedings, volume 10843 of Lecture Notes in Computer Science, pp. 593–607. Springer, 2018.
  Min Joon Seo, Aniruddha Kembhavi, Ali Farhadi, and Hannaneh Hajishirzi. Bidirectional attention flow for machine comprehension. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
  David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nat., 529(7587):484–489, 2016.
  David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140–1144, 2018.
  Ilya Sutskever, Oriol Vinyals, and Quoc V. Le. Sequence to sequence learning with neural networks. In Zoubin Ghahramani, Max Welling, Corinna Cortes, Neil D. Lawrence, and Kilian Q. Weinberger (eds.), Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada, pp. 3104–3112, 2014.
  Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
  Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Information Theory Workshop, ITW 2015, Jerusalem, Israel, April 26 - May 1, 2015, pp. 1–5. IEEE, 2015.
  Adam Trischler, Marc-Alexandre Côté, and Pedro Lima. First TextWorld Problems, the competition: Using text-based games to advance capabilities of AI agents. 2019.
  Théo Trouillon, Johannes Welbl, Sebastian Riedel, Éric Gaussier, and Guillaume Bouchard. Complex embeddings for simple link prediction. In Maria-Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 2071–2080. JMLR.org, 2016.
  Jack Urbanek, Angela Fan, Siddharth Karamcheti, Saachi Jain, Samuel Humeau, Emily Dinan, Tim Rocktäschel, Douwe Kiela, Arthur Szlam, and Jason Weston. Learning to speak and act in a fantasy text adventure game. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLP 2019, Hong Kong, China, November 3-7, 2019, pp. 673–683. Association for Computational Linguistics, 2019.
  Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of machine learning research, 9(11), 2008.
  Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998–6008, 2017.
  Arthur Wandzel, Yoonseon Oh, Michael Fishman, Nishanth Kumar, Lawson L. S. Wong, and Stefanie Tellex. Multi-object search using object-oriented pomdps. In International Conference on Robotics and Automation, ICRA 2019, Montreal, QC, Canada, May 20-24, 2019, pp. 7194–7200. IEEE, 2019.
  Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforcement learning. CoRR, abs/1907.02057, 2019.
  Xusen Yin and Jonathan May. Comprehensible context-driven text game playing. In IEEE Conference on Games, CoG 2019, London, United Kingdom, August 20-23, 2019, pp. 1–8. IEEE, 2019.
  Xingdi Yuan, Marc-Alexandre Côté, Alessandro Sordoni, Romain Laroche, Remi Tachet des Combes, Matthew J. Hausknecht, and Adam Trischler. Counting to explore and generalize in text-based games. CoRR, abs/1806.11525, 2018.
  Tom Zahavy, Matan Haroush, Nadav Merlis, Daniel J. Mankowitz, and Shie Mannor. Learn what not to learn: Action elimination with deep reinforcement learning. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 3566–3577, 2018.
  Mikuláŝ Zelinka. Baselines for reinforcement learning in text games. In Lefteri H. Tsoukalas, Éric Grégoire, and Miltiadis Alamaniotis (eds.), IEEE 30th International Conference on Tools with Artificial Intelligence, ICTAI 2018, 5-7 November 2018, Volos, Greece, pp. 320–327. IEEE, 2018.
  Guangxiang Zhu, Zhiao Huang, and Chongjie Zhang. Object-oriented dynamics predictor. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 9826–9837, 2018.
  Guangxiang Zhu, Jianhao Wang, Zhizhou Ren, Zichuan Lin, and Chongjie Zhang. Object-oriented dynamics learning through multi-level abstraction. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 6989–6998. AAAI Press, 2020.
sections.heading: A MODEL IMPLEMENTATION DETAILS AND DATASETS
sections.heading: A.1 HYPER-PARAMETERS
sections.text: We report the key hyper-parameters that have been applied for model implementation. 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. The learning rate of policy training and dynamics model training is set to 0.0001. The discount factor γ is set to 0.9. The size of the hidden layers and the size of z in our OOTD model are set to 32. 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. The sizes of word embedding and node embedding are set to 300 and 100 respectively. These numbers are determined experimentally, which could provide promising performance with a reasonable computational resource. The full list of hyper-parameters is recorded in the .yaml file in our submitted code. In our experiment, the random seeds we select are 123, 321, and 666.
sections.heading: A.2 RESAMPLING FOR COMPUTING THE GRAPH LIKELIHOOD
sections.text: The graph likelihood pΩ(ht|zt) computes the probability of generating ht with a graph decoder. As it is discussed in Section 3.1, our graph decoder is implemented as a scoring function for entry-level prediction. However, ht is a graph with significant 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. As a result, directly optimizing toward h∗t creates a large negative bias. In other words, the decoder can achieve a satisfying accuracy by labelling all the entries with 0s, which does not capture any object information. 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. These positive and negative samples are used to estimate the graph likelihood.
sections.heading: A.3 GOAL EXTRACTOR
sections.text: At a time step t, the goal extractor generates the goal gt from the object states zt by optionally conditioning on the previous goal gt−1 and observed text ot. The goal can be represented by a sequence of words and symbols that follows the triplet format. 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.
  Goal Extractor with the OS-ELBo objective. 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). We can build a rule-based goal extractor that maps zt to ĥOSt and checks whether hOSt captures certain relations (e.g., ’need’) between two objects. If there is such a relation, we directly output the sentence of corresponding triplets by following a format "[ object 1 $ relation $ object 2 ]". Such a simple extractor can generalize well to training and testing games if ĥOSt manages to capture correct relations.
  Goal Extractor with the SS-ELBo objective. The memory graph ĥOSt learned by the SS-ELBo objective captures a latent representation of ĥ∗t . 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 (gt) and inputs are zt, gt−1 and ot. To train this model, we manually build a goal learning dataset that collects pairs of model input/output from the training games. During testing, we predict the goal after interacting with the environment and update the current goal if a valid goal sentence is generated.
sections.heading: A.4 OBJECT EXTRACTOR
sections.text: The object extractor fe(ot, at−1, ht−1) predicts the memory graph ht by conditioning on current observations ot and action at. Figure 5 shows the architecture of fe. It shares a similar structure to our OOTD model, but the key differences are 1) fe is a deterministic model which only maximizes the likelihood of predicting h∗t during training. 2) fe captures object information from ot, whereas
  ot is not available for our OOTD model, and it is the motivation why OOTD predicts the belief (distribution) of states while fe directly generates the deterministic states.
sections.heading: A.5 OBSERVATION ENCODER.
sections.text: The observation encoder qE is defined by: νak,t−1 = ∑
  j bak,jψ a(aj,t−1) where bak = softmax(B ae k,:) ∈ [0, 1]J , νet−1 = ∑
  k bekek,t−1 where b e = softmax(maxcol(B ae)) ∈ [0, 1]K , νok,t−1 = ∑
  j bok,jψ o(aj,t−1) where bok = softmax(B oe k,:) ∈ [0, 1]J ν′,et−1 = ∑
  k boe,ek ek,t−1 where b oe,e = softmax(maxcol(B oe)) ∈ [0, 1]K (8)
  qkE(zk,t|ot,at−1, zk,t−1) = N (µ q k,t,σ q k,t) where [µ q k,t,σ q k,t] = ψ q k([ν o k,t−1,ν ′,e t−1,ν a k,t−1,ν e t−1, ek,t−1])
  ψo is a transformer-based text embedding function (Vaswani et al., 2017) for observations, and ψqk is implemented by a MLP. The transition layers ψq1, . . . , ψ q K are parameterized independently.
  Figure 6 shows the structure of an observation encoder qkE , which shares a similar structure to the transition (OOTD) model (Figure 3), but it includes an additional BIDAF network to learn an attention matrixBoe, where boek,i measures the similarity between the representation of the k
  th object and the ith word in the current observation ot. It predicts the memory graph ht from z q t with a ComplEx graph decoder (Section 3.1). It predicts the current observation ôt and reward rt with a sequence decoder and the reward function (Section 3.2).
sections.heading: Post MLP
sections.heading: A.6 FTWP DATASET
sections.text: The FTWP dataset is a public dataset that supports pre-training the dynamics model 1. Trischler et al. (2019) provided the First TextWorld Problems (FTWP) dataset. The dataset consists of TextWorld
  1The dataset can be downloaded at https://aka.ms/ftwp/dataset.zip
  games in a cooking theme across a wide range of difficulty levels. For each game, the expert trajectory and the memory graph ht at each time step are recorded. The original dataset has only 10 games for each difficulty level. This is insufficient for dynamics model pre-training, so Adhikari et al. (2020) further expanded this dataset by adding more games into this dataset by utilizing the logic in the Textworld game engine. The specific generation rules are described in (Adhikari et al., 2020). After expansion, the dataset contains 4440 games for training, 222 games for validation, 514 games for testing. To ensure fairness when using this dataset, we confirm that there is no overlap between the FTWP and the games we used to train and evaluate our planning algorithms.
sections.heading: A.7 COMPUTATIONAL RESOURCE AND RUNNING TIME
sections.text: We run the experiment on a cluster operated by the Slurm workload manager. 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. We used machines with 24 GB of memory for pre-training the object extractor and 64 GB for training the OOTD model. The number of running nodes is 1, and the number of CPUs requested per task is 32. Given the aforementioned resources, pre-training can be completed within 4 days (96 hours), and training the dynamics model takes 6 days (144 hours).
sections.heading: 4 3 6 11 X X 28.4 33.4
sections.heading: 3 1 9 3 7 7 7.2 34.1
sections.heading: 2 1 1 5 X X 11.8 17.5
sections.heading: 1 1 1 4 X 7 11.5 17.1
sections.heading: B ALGORITHMS
sections.text: Algorithm 1: Dyna-Q Training Input :Transition Model: pT (zt+1|at, zt),
  Object Encoder: qE(zt+1|ot+1, at, zt), Reward Model: pr(rt+1|zt+1, at, gt), Object Extraction function: fe(ht, at, ot+1), Graph Decoder: pΩ(ht|zt), Training game ids I.
  Initialize Replay Buffers D; while not converged do
  Initialize belief state z0 =random numbers, action a0 ="restart", goal g0 ="<pad>" and memory graph h0 = 0C×K×K ; Randomly pick a game id ∈ I, set t = 0; while not end do
  /* Q learning with environment */ ot+1, rt+1 ← env(at, id); Extract object information ht+1 = f(at, zt, ot+1); Update object belief zt+1 ∼ qE(zt+1|ot+1, at, zt); Extract goal: gt+1 = ExtactGoal(zt+1, ot+1, gt); Action selection: ât+1 = EpsilonGreedy(Q(zt+1, at+1, gt+1)); Expand the replay buffer: D ∪ {zt, ht, at, gt, rt+1, zt+1, ht+1, gt+1}; t = t+ 1 and end = CheckEnd(env)
  end for update step c=1,...,C do
  /* Model fitting */ Sample a transition: T = {zt, ht, at, gt, rt+1, zt+1, ht+1, gt+1} from the dataset D; Update the reward model to minimize the L1Smooth loss (Equation 6) with T ; Update the transition model to minimize the OS-ELBO (Equation 5) objective or the
  SS-ELBO objective (Equation 7) with T ; 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 */ Randomly select a′τ ∈ A ; z′τ+1 ∼ pT (z′τ+1|a′τ , zτ ), r′τ+1 ∼ pr(r′t+1|z′τ+1, a′τ , gτ ); Update the transition T ′ = {zt, gt, a′t, r′t+1, z′t+1, gt+1}; Update the Q function to minimize the TD Loss (Equation 1) with T ;
  end end
  end
  Algorithm 2: MCTS Testing Input :Transition Model: pT (zt+1|at, zt),
  Object Encoder: qE(zt+1|ot+1, at, zt), Reward Model: pr(rt+1|zt+1, at, gt), Action Value function: Q(zt, at, gt), Testing game ids I
  Initialize belief state z0 =random numbers, action a0 ="restart", and goal g0 ="<pad>"; Initialize reward bufferR; for game id id ∈ I do
  Initialize root N0 = InitTree(Q(zt, at, gt)), set t = 0 and R = 0; while not end do
  ot+1, rt+1 ← env(at, id) and R = R+ rt+1; Update object belief zt+1 ∼ qE(zt+1|ot+1, at, zt); Extract goal: gt+1 = ExtactGoal(zt+1, ot+1, gt); Select action: ât+1 = MCTS_Simulations(Nt, zt+1, gt+1, pT (·), pr(·), Q(·)); Move the search node: Nt+1 = Move(Nt, ât+1); Expand t = t+ 1 and end = CheckEnd(env)
  end Buffer total rewardsR = R∪R
  end output :Reward bufferR
sections.heading: C COMPLEMENTARY RESULTS
sections.heading: C.1 GRAPH VISUALIZATION
sections.text: Graph Visualization. Figure 7 illustrates examples of the ground-truth memory graph h∗t and the graphs learned by the OS-ELBo (hOSt ) and SS-ELBo (h SS t ) objectives. 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. Supervised by predicted graphs ĥt, OS-ELBo captures the sparsity, and hOSt is a direct approximation to h∗t . However, graph information are unavailable for the SS-ELBo objective, so hOSt forms a dense latent representation of h ∗ t . Similar phenomenon can be observed in GATA (Adhikari et al., 2020). We study how well this latent representation and the corresponding object states can represent the object dynamics in Section 4.2.
sections.heading: C.2 PLANNING EFFICIENCY
sections.text: We compare the efficiency 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. Figure 8 shows results. Among the planning algorithms, Dyna-Q+MCTS significantly outperforms others. 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. Together, they yield a more accurate look ahead of future rewards.