Learning Object-Oriented Dynamics for Planning from Text

1 INTRODUCTION . 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 . 2 OBJECT LEVEL PLANNING IN TEXT DOMAINS . 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 . 2.1 TEXT-BASED GAMES . 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 . 2.2 OBJECT-ORIENTED PARTIALLY OBSERVABLE MARKOV DECISION PROCESSES . 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 ) . 2.3 GRAPH REPRESENTATION FOR OBJECTS . 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 can not 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 . 2.4 MODEL-BASED PLANNING IN TEXT-BASED GAMES . 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 ) .

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