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