SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yue Wu $^{1*}$ , Yewen Fan $^{1}$ , Paul Pu Liang $^{1}$ , Amos Azaria $^{2}$ , Yanzhi Li $^{1,3}$ , Tom Mitchell $^{1}$ $^{1}$ Carnegie Mellon University,  $^{2}$ Ariel University,  $^{3}$ Microsoft Research ywu5@andrew.cmu.edu
High sample complexity has long been a challenge for RL. On the other hand, humans learn to perform tasks not only from interaction or demonstrations, but also by reading unstructured text documents, e.g., instruction manuals. Instruction manuals and wiki pages are among the most abundant data that could inform agents of valuable features and policies or task-specific environmental dynamics and reward structures. Therefore, we hypothesize that the ability to utilize human-written instruction manuals to assist learning policies for specific tasks should lead to a more efficient and better-performing agent. We propose the Read and Reward framework. Read and Reward speeds up RL algorithms on Atari games by reading manuals released by the Atari game developers. Our framework consists of a QA Extraction module that extracts and summarizes relevant information from the manual and a Reasoning module that evaluates object-agent interactions based on information from the manual. An auxiliary reward is then provided to a standard A2C RL agent, when interaction is detected. Experimentally, various RL algorithms obtain significant improvement in performance and training speed when assisted by our design. Code at github.com/Holmeswww/RnR.
Reinforcement Learning (RL) has achieved impressive performance in a variety of tasks, such as Atari (Mnih et al., 2015; Badia et al., 2020), Go (Silver et al., 2017), or autonomous driving (Fuchs et al., 2021; Kiran et al., 2022), and is hypothesized to be an important step toward artificial intelligence (Silver et al., 2021). However, RL still faces a great challenge when applied to complicated, real-world-like scenarios (Shridhar et al., 2020b; Szot et al., 2021) due to its low sample efficiency. The Skiing game in Atari, for example, requires the skier (agent) to ski down a snowy hill and hit the objective gates in the shortest time, while avoiding obstacles. Such an intuitive game with simple controls (left, right) still required 80 billion frames to solve with existing RL algorithms (Badia et al., 2020), roughly equivalent to 100 years of non-stop playing.
Observing the gap between RL and human performance, we identify an important cause: the lack of knowledge and understanding about the game. One of the most available sources of information, that is often used by humans to solve real-world problems, is unstructured text, e.g., books, instruction manuals, and wiki pages. In addition, unstructured language has been experimentally verified as a powerful and flexible source of knowledge for complex tasks (Tessler et al., 2021). Therefore, we hypothesize that RL agents could benefit from reading and understanding publicly available instruction manuals and wiki pages.
Figure 1: An overview of our Read and Reward framework. Our system receives the current frame in the environment, and the instruction manual as input. After object detection and grounding, the QA Extraction Module extracts and summarizes relevant information from the manual, and the Reasoning Module assigns auxiliary rewards to detected in-game events by reasoning with outputs from the QA Extraction Module. The “Yes/No” answers are then mapped to  $+5 / -5$  auxiliary rewards.
In this work, we demonstrate the possibility for RL agents to benefit from human-written instruction manuals. We download per-game instruction manuals released by the original game developers<sup>1</sup> or Wikipedia and experiment with 4 Atari games where the object location ground-truths are available.
Two challenges arise with the use of human-written instruction manuals: 1) Length: Manuals contain a lot of redundant information and are too long for current language models. However, only a few paragraphs out of the 2-3 page manuals contain relevant information for any specific aspect of concern, i.e., the game's objective or a specific interaction. 2) Reasoning: References may be implicit and attribution may require reasoning across paragraphs. For example, in Skiing, the agent needs to understand that hitting an obstacle is bad using both "you lose time if the skier hits an obstacle" and "the objective is to finish in shortest time", which are from different parts of the manual. Furthermore, the grounding of in-game objects to references in the text remains an open problem and manuals may contain significant syntactic variation. For example, the phrase "don't hit an obstacle" is also referred to as "don't get wrapped around a tree" or "avoid obstacles". Finally, the grounding of in-game features to references in the text remains an open problem (Luketina et al., 2019).
Our proposed Read and Reward framework (Figure 1) addresses the challenges with a two-step process. First, a zero-shot extractive QA module is used to extract and summarize relevant information from the manual, and thus tackles the challenge of Length. Second, a zero-shot reasoning module, powered by a large language model, reasons about contexts provided by the extractive QA module and assigns auxiliary rewards to in-game interactions. After object detection and localization, our system registers an interaction based on the distance between objects. The auxiliary reward from the interaction can be consumed by any RL algorithm.
Even though we only consider "hit" interactions (detected by tracking distance between objects) in our experiments, Read and Reward still achieves  $60\%$  performance improvement on a baseline using 1000x fewer frames than the SOTA for the game of Skiing. In addition, since Read and Reward does not require explicit training on synthetic text (Hermann et al., 2017; Chaplot et al., 2018; Janner et al., 2018; Narasimhan et al., 2018; Zhong et al., 2019, 2021; Wang & Narasimhan, 2021), we observe consistent performance on manuals from two different sources: Wikipedia and Official manuals.
To our knowledge, our work is the first to demonstrate the capability to improve RL performance in an end-to-end setting using real-world manuals designed for human readers.
$^{1}$ atariage.com
2.1 Reducing sample complexity in RL
Prior works have attempted to improve sample efficiency with new exploration techniques (Schulman et al., 2017; Haarnoja et al., 2018; Badia et al., 2020), self-supervised objectives (Pathak et al., 2017; Schrittwieser et al., 2020), and static demonstration data (Kumar et al., 2019, 2020). Fan & Xiao (2022) cast RL training into a training data distribution optimization problem, and propose a policy mathematically controlled to traverse high-value and non-trivial states. However, such a solution requires manual reward shaping. Currently, the SOTA RL agent (Badia et al., 2020) still spends billions of frames to learn the Skiing game that can be mastered by humans within minutes.
2.2 Grounding objects for control problems
Most prior works study grounding – associating perceived objects with appropriate text – in the setting of step-by-step instruction following for object manipulation tasks (Wang et al., 2016; Bahdanau et al., 2018) or indoor navigation tasks (Chaplot et al., 2018; Janner et al., 2018; Chen et al., 2019; Shridhar et al., 2020a). The common approach has been to condition the agent's policy on the embedding of both the instruction and observation (Mei et al., 2016; Hermann et al., 2017; Janner et al., 2018; Misra et al., 2017; Chen et al., 2019; Pashevich et al., 2021). Recently, modular solutions have been shown to be promising (Min et al., 2021). However, the majority of these works use synthetic language. These often take the form of templates, e.g., “what colour is <object> in <room>” (Chevalier-Boisvert et al., 2018). On the other hand, recent attempts on large-scale image-text backbones, e.g., CLIP (Radford et al., 2021) have shown robustness at different styles of visual inputs, even on sketches. Therefore, we pick CLIP as our model for zero-shot grounding.
2.3 Reinforcement learning informed by natural language
In the language-assisted setting, step-by-step instructions have been used to generate auxiliary rewards, when environment rewards are sparse. Goyal et al. (2019); Wang et al. (2019) use auxiliary reward-learning modules trained offline to predict whether trajectory segments correspond to natural language annotations of expert trajectories.
In a more general setting, text may contain both information about optimal policy and environment dynamics. Branavan et al. (2012) improve Monte-Carlo tree search planning by accessing a natural language manual. They apply their approach to the first few moves of the game Civilization II, a turn-based strategy game. However, many of the features they use are handcrafted to fit the game of Civilization II, which limits the generalization of their approach.
Narasimhan et al. (2018); Wang & Narasimhan (2021) investigate planning in a 2D game environment with a fixed number of entities that are annotated by natural language (e.g. the 'spider' and 'scorpion' entities might be annotated with the descriptions "randomly moving enemy" and "an enemy who chases you", respectively). The agent learns to generalize to new environment settings by learning the correlation between the annotations and the environmental goals and mechanics. The proposed agent achieves better performance than the baselines, which do not use natural language. However the design of the 2D environment uses 6-line descriptions created from templates. Such design oversimplifies the observations into a 2D integer matrix, and removes the need for visual understanding and generalization to new objects and new formats of instruction manual, e.g., human written manuals.
2.4 RL Models that Read Natural Language Instructions
Zhong et al. (2019, 2021) make use of special architectures with Bidirectional Feature-wise Linear Modulation Layers that support multi-hop reasoning on  $6 \times 6$  grid worlds with template-generated instruction manuals. However, the model requires 200 million training samples from templates identical to the test environments. Such a training requirement results in performance loss even on  $10 \times 10$  grid worlds with identical mechanics, thus limiting the generalization of the model.
Wang & Narasimhan (2021) mixes embedding of entities with multi-modal attention to get entity representations, which are then fed to a CNN actor. The attention model and CNN actor are trained jointly using RL. However, the whole framework is designed primarily for a grid world with a template generated instruction manual consisting of one line per entity. In our experiments, we find
Figure 2: Illustration of the QA Extraction Module on the game PacMan. We obtain generic information about the game by running extractive QA on 4 generic questions (3 shown since one question did not have an answer). We then obtain object-specific information using a question template. We concatenate generic and object-specific information to obtain the <context> string.
that extracting entity embeddings from human-written Atari instruction manuals following Wang & Narasimhan (2021) does not improve the performance of our vision-based agent due to the complexity and diversity of the language.
3 Read and Reward
We identify two key challenges to understanding and using the instruction manual with an RL agent.
The first challenge is Length. Due to the sheer length of the raw text (2 ~ 3 pages), current pre-trained language models cannot handle the raw manual due to input length constraints. Additionally, most works in the area of long-document reading and comprehension (Beltagy et al., 2020; Ainslie et al., 2020; Zemlyanskiy et al., 2021) require a fine-tuning signal from the task, which is impractical for RL problems since the RL loss already suffers from high variance. Furthermore, an instruction manual may contain a lot of task-irrelevant information. For example, the official manual for MsPacman begins with: "MS. PAC-MAN and characters are trademark of Bally Midway Mfg. Co. sublicense to Atari, Inc. by Namco-America, Inc."
It is therefore intuitive in our first step to have a QA Extraction Module that produces a summary directed toward game features and objectives, to remove distractions and simplify the problem of reasoning.
The second challenge is Reasoning. Information in the manual may be implicit. For example, the Skiing manual states that the goal is to arrive at the end in the fastest time, and that the player will lose time when they hit a tree, but it never states directly that hitting a tree reduces the final reward. Most prior works (Eisenstein et al., 2009; Narasimhan et al., 2018; Wang & Narasimhan, 2021; Zhong et al., 2019, 2021) either lack the capability of multi-hop reasoning or require extensive training to form reasoning for specific manual formats, limiting the generalization to real-world manuals with a lot of variations like the Atari manuals.
On the other hand, large language models have achieved success without training in a lot of fields including reading comprehension, reasoning, and planning (Brown et al., 2020; Kojima et al., 2022; Ahn et al., 2022). Therefore, in the Reasoning Module, we compose natural language queries about specific object interactions in the games, to borrow the in-context reasoning power of large language models.
The rest of the section describes how the two modules within our Read and Reward framework (a QA extraction module and a reasoning module) are implemented.
QA Extraction Module (Read) In order to produce a summary of the manual directed towards game features and objectives, we follow the extractive QA framework first proposed by Devlin et al. (2018), which extracts a text sub-sequence as the answer to a question about a passage of text. The extractive QA framework takes raw text sequence  $S_{\text{manual}} = \{w^0, \dots, w^L\}$  and a question string  $S_{\text{question}}$  as input. Then for each token (word)  $w^i$  in  $S_{\text{manual}}$ , the model predicts the probability that the current token is the start token  $p_{\text{start}}^i$  and end token  $p_{\text{end}}^i$  with a linear layer on top of word piece embeddings. Finally, the output  $S_{\text{answer}} = \{w^{\text{start}}, \dots, w^{\text{end}}\}$  is selected as a sub-sequence of  $S_{\text{manual}}$  that maximizes the overall probability of the start and end:  $w^{\text{start}}, w^{\text{end}} = \arg \max_{w^{\text{start}}, w^{\text{end}}} p^{\text{start}} p^{\text{end}}$ .
Implementation-wise, we use a RoBERTa-large model (Liu et al., 2019) for Extractive QA fine-tuned on the SQUAD dataset (Rajpurkar et al., 2016). Model weights are available from the AllenNLP  $\mathrm{API}^2$ . To handle instruction manual inputs longer than the maximum length accepted by RoBERTa, we split the manual into chunks according to the max-token size of the model, and concatenate extractive QA outputs on each chunk.
As shown in Figure 2, we compose a set of generic prompts about the objective of the game, for example, "What is the objective of the game?" In addition, to capture information on agent-object interactions, we first identify the top 10 important objects using TFIDF. Then for each object, we obtain an answer to the generated query: "What happens when the player hit a  $<$  object>?"
Finally, we concatenate all non-empty question-answer pairs per interaction into a <context> string. Note that the format of context strings closely resembles that of a QA dialogue.
Zero-shot reasoning with pre-trained QA model (Reward) The summary <context> string from the QA Extraction module could directly be used as prompts for reasoning through large language models (Brown et al., 2020; Kojima et al., 2022; Ahn et al., 2022).
Motivated by the dialogue structure of the summary, and the limit in computational resources, we choose Macaw (Tafjord & Clark, 2021), a general-purpose zero-shot QA model with performance comparable to GPT-3 (Brown et al., 2020). Compared to the RoBERTa extractive QA model for the QA Extraction module, Macaw is significantly larger and is better suited for reasoning. Although we find that GPT-3 generally provides more flexibility and a better explanation for its answers for our task, Macaw is faster and open-source.
As shown in Figure 3, for each interaction we compose the query prompt as "<context> Question: Should you hit a <object of interaction> if you want to win? Answer: ", and calculate the LLM score on choices  $\{\mathrm{Yes},\mathrm{No}\}$  .The Macaw
model directly supports scoring the options in a multiple-choice setting. We note that for other generative LLMs, we can use the probability of "Yes" vs. "No" as scores, following Ahn et al. (2022).
Finally, during game-play, a rule-based algorithm detects 'hit' events by checking the distance between bounding boxes. For detected interactions, an auxiliary reward is provided to the RL algorithm according to  $\{\mathrm{Yes},\mathrm{No}\}$  rating from the reasoning module (see Section 4.4 for details).
4.1 Atari Environment and Baselines
The Openai gym Atari environment contains diverse Atari games designed to pose a challenge for human players. The observation consists of a single game screen (frame): a 2D array of 7-bit pixels, 160 pixels wide by 210 pixels high. The action space consists of the 18 discrete actions defined by the joystick controller. Instruction manuals released by the original game developers have been scanned and parsed into publicly available HTML format<sup>3</sup>. For completeness, we also attach an explanation of the objectives for each game in Appendix A.3.
Figure 3: Illustration of the Reasoning Module. The <context> (from Figure 2) related to the object ghost from the QA Extraction module is concatenated with a template-generated question to form the zero-shot in-context reasoning prompt for a Large Language Model. The Yes/No answer from the LLM is then turned into an auxiliary reward for the agent.
2AllenNLP Transformer-QA
3 atariage.com
4.2 Delayed Reward Schedule
Most current Atari environments include a dense reward structure, i.e., the reward is obtained quite often from the environment. Indeed, most RL algorithms for the Atari environments perform well due to their dense reward structure. However, in many real-world scenarios (Kolve et al., 2017; Wang et al., 2022a) or open-world environments (Fan et al., 2022), it is expensive to provide a dense reward. The reward is usually limited and "delayed" to a single positive or negative reward obtained at the very end of an episode (e.g., once the robot succeeds or fails to pour the coffee).
Notably, Atari games with the above "delayed" reward structure, such as Skiing, pose great challenges for current RL algorithms, and are among some of the hardest games to solve in Atari (Badia et al., 2020). Therefore, to better align with real-world environments, we use a delayed reward version of the tennis, breakout, and Pac-Man games by providing the reward (the "final game score") only at the end of the game. This delayed reward structure does not change the optimization problem, but is more realistic and imposes additional challenges to RL algorithms.
4.3 Grounding Objects in Atari
We find grounding (detecting and relating visual objects to keywords from TFIDF) a significant challenge to applying our framework. In some extensively studied planning/control problems, such as indoor navigation, pre-trained models such as masked RCNN reasonably solve the task of grounding (Shridhar et al., 2020a). However, the objects in the Atari environments are too different from the training distribution for visual backbones like mask RCNN. Therefore, the problem becomes very challenging for the domain of Atari. Since it is not the main focus of our proposed work, we attempt to provide an unsupervised end-to-end solution only for the game of Skiing.
Full end-to-end pipeline on Skiing with A2C
To demonstrate the full potential of our work, we offer a proof-of-concept agent directly operating on raw visual game observations and the instruction manual for the game of Skiing. For unsupervised object localization, we use SPACE, a method that uses spatial attention and decomposition to detect visual objects (Lin et al., 2020). SPACE is trained on 50000 random observations from the environment, and generates object bounding boxes. These bounding boxes are fed to CLIP, a zero-shot model that can pair images with natural language descriptions (Radford et al., 2021), to classify the bounding box
into categories defined by key-words from TF-IDF. The SPACE/CLIP models produce mostly reasonable results; however, in practice we find that CLIP lacks reliability over time, and SPACE cannot distinguish objects that cover each other. Therefore, to improve classification reliability, we use a Multiple Object Tracker (Bewley et al., 2016) in conjunction with CLIP and set the class of the bounding box to be the most dominant class over time.
Ground-truth object localization/grounding experiments For other games (Tennis, Ms. Pac-Man, and Breakout), we follow Anand et al. (2019) and directly calculate the coordinates of all game objects from the simulator RAM state. Specifically, we first manually label around 15 frames with object name and X,Y coordinates and then train a linear regression model mapping from simulator RAM state to labeled object coordinates. Note that this labeling strategy fails for games in which multiple objects of the same kind appear, i.e., multiple trees appear in one frame of Skiing.
Figure 4: Examples of SPACE (Lin et al., 2020) and CLIP (Radford et al., 2021) in the full end-to-end pipeline (Section 4.3). The top row shows bounding boxes for objects and the bottom row shows corresponding object masks as detected by SPACE. Most of the bounding boxes generated are correct. Left: SPACE confuses bounding boxes of agent and tree into one and the box gets classified as "tree" (blue), and the auxiliary penalty is not properly triggered. Right: The flag next to the agent (in red circle) is not detected, and therefore the auxiliary reward is not provided.
A2C+R&R	A2C	PPO+R&R	PPO	R2D1+R&R	R2D1	Agent57+R&R	Agent57
Tennis	-5 (-5)	-23	-6 (-6)	-23	-1 (-1)	-23	0 (0)	-23
Pacman	580 (580)	452	253 (253)	204	3115 (3115)	2116	1000 (1000)	708
Breakout	14 (14)	2	9 (9)	2	10 (10)	2	38 (38)	2
Table 1: Table of the game score of different RL algorithms trained under delayed reward schedule, RL with Read and Reward consistently outperforms their counterparts that are not using the manual. In addition, Read and Reward performance remains consistent across official Atari manual and Wikipedia in brackets “(.”).
A2C+R&R	A2C	PPO+R&R	PPO	R2D1+R&R	R2D1
Tennis	-5.2 ± 2.2	-23 ± 1.3	-8.1 ± 2.3	-23 ± 1.0	-2.2 ± 3.1	-23.0 ± 1.3
Pacman	455.2 ± 63.2	387.1 ± 66.2	-284.2 ± 67.3	200.1 ± 65.4	3001.3 ± 203.2	1999.2 ± 109.2
Breakout	-4.5 ± 3.0	2.1 ± 0.3	-10.2 ± 0.1	1.9 ± 0.5	10 ± 3.0	2.1 ± 0.2
Table 2: Table of the game score of different RL algorithms trained under delayed reward schedule on the official Atari manual, with statistics computed over 3 random seeds. RL with Read and Reward consistently outperforms their counterparts that are not using the manual.
4.4 Simplifying Interaction Detection to Distance Tracking
We make the assumption that all object interactions in Atari can be characterized by the distance between objects, and we ignore all interactions that do not contain the agent. We note that since the Atari environment is in 2D, it suffices to consider only interactions where objects get close to the agent so that we can query the manual for "hit" interaction.
We, therefore, monitor the environment frames and register interaction whenever an object's bounding boxes intersect with the agent. For each detected interaction, we assign rewards from  $\{-r_n,r_p\}$  accordingly as detailed in Section 3. Experimentally, we find that it suffices to set  $r_n = r_p = 5$
4.5 RL Baselines