[Paper Review] Learning Dynamic Belief Graphs to Generalize on Text-Based Games

Ashutosh Adhikari, Xingdi Yuan+8 more|arXiv (Cornell University)|Feb 21, 2020

Topic ModelingComputer Science56 references55 citations

TL;DR

The paper introduces GATA, a graph-aided transformer agent that learns to build and update latent belief graphs from text observations to plan and generalize across many TextWorld text-based games, outperforming text-based baselines.

ABSTRACT

Playing text-based games requires skills in processing natural language and sequential decision making. Achieving human-level performance on text-based games remains an open challenge, and prior research has largely relied on hand-crafted structured representations and heuristics. In this work, we investigate how an agent can plan and generalize in text-based games using graph-structured representations learned end-to-end from raw text. We propose a novel graph-aided transformer agent (GATA) that infers and updates latent belief graphs during planning to enable effective action selection by capturing the underlying game dynamics. GATA is trained using a combination of reinforcement and self-supervised learning. Our work demonstrates that the learned graph-based representations help agents converge to better policies than their text-only counterparts and facilitate effective generalization across game configurations. Experiments on 500+ unique games from the TextWorld suite show that our best agent outperforms text-based baselines by an average of 24.2%.

Motivation & Objective

• Motivate solving text-based games with learned graph-structured state representations instead of hand-crafted heuristics.
• Develop a graph-augmented transformer agent (GATA) that infers and updates latent belief graphs during planning.
• Pre-train the graph updater with self-supervised tasks to capture environment dynamics.
• Train the action selector with reinforcement learning to optimize rewards.
• Evaluate generalization across a distribution of unseen TextWorld games and compare to ground-truth graph baselines.

Proposed method

• Represent environment state as a dynamic belief graph G; G is updated as a latent adjacency tensor.
• Use a graph updater to predict graph changes Delta g_t from previous graph, observation, and last action; G_t = G_{t-1} ⊕ Delta g_t.
• Pre-train graph updater with Observation Generation (Seq2Seq reconstruction of O_t from G_t and A_{t-1}) and Contrastive Observation Classification (maximize MI between G_t and O_t).
• Encode G_t with Relational Graph Convolutional Networks (R-GCN) and adapt relation embeddings; encode O_t with a Transformer; fuse via bi-directional attention for action scoring.
• Train the action selector with Double DQN (multi-step, prioritized replay) across episodes drawn from multiple games to enable cross-game generalization.

Experimental results

Research questions

• RQ1 Can a learned, dynamic, graph-structured representation improve policy learning and generalization in text-based games?
• RQ2 How close can a learned belief graph-based agent approach the performance of agents with access to ground-truth graphs?
• RQ3 Do self-supervised pre-training tasks help the graph updater encode useful environment dynamics from partial text observations?

Key findings

• GATA outperforms strong text-based baselines (including transformer-equipped DQN variants) on unseen TextWorld games.
• Pre-training the graph updater with both self-supervised tasks yields better performance than using either alone.
• Using text observations alongside the belief graph further improves policy quality via attention-guided scoring.
• Agents with access to ground-truth graphs (GATA-GTF) achieve higher scores than GATA but still demonstrate the benefit of learned graphs in partially observed settings.
• GATA-GTF represents an upper bound for GATA-style graph-based methods, highlighting improved state representations as a key factor for performance.
• Across difficulty levels, GATA demonstrates consistent improvements in policy quality and generalization over text-only baselines.