Contrastive Reinforcement Learning of Symbolic Reasoning Domains

Contrastive Reinforcement Learning of Symbolic Reasoning Domains Gabriel Poesia Stanford University poesia @ cs.stanford.edu WenXin Dong Stanford University wxd @ stanford.edu Noah Goodman Stanford University ngoodman @ stanford.edu Abstract Abstract symbolic reasoning , as required in domains such as mathematics and logic , is a key component of human intelligence . Solvers for these domains have important applications , especially to computer-assisted education . But learning to solve symbolic problems is challenging for machine learning algorithms . Existing models either learn from human solutions or use hand-engineered features , making them expensive to apply in new domains . In this paper , we instead consider symbolic domains as simple environments where states and actions are given as unstructured text , and binary rewards indicate whether a problem is solved . This flexible setup makes it easy to specify new domains , but search and planning become challenging . We introduce five environments inspired by the Mathematics Common Core Curriculum , and observe that existing Reinforcement Learning baselines perform poorly . We then present a novel learning algorithm , Contrastive Policy Learning ( ConPoLe ) that explicitly optimizes the InfoNCE loss , which lower bounds the mutual information between the current state and next states that continue on a path to the solution . ConPoLe successfully solves all four domains . Moreover , problem representations learned by ConPoLe enable accurate prediction of the categories of problems in a real mathematics curriculum . Our results suggest new directions for reinforcement learning in symbolic domains , as well as applications to mathematics education . 1 Introduction . Humans posses the remarkable ability to learn how to reason in symbolic domains , such as arithmetic , algebra , and formal logic . Our aptitude for mathematical cognition builds on specialized neural bases but extends them radically through formal education [ 6 , 16 , 11 , 12 ] . Learning to reason in symbolic domains poses an important challenge for artificial intelligence research . As we describe below , this type of reasoning has unique features that distinguish it from domains in which machine learning has had recent success . From a practical viewpoint , since symbolic reasoning skills span years of instruction in school , advances in symbolic reasoning may have a large impact on education . In particular , automated tutors equipped with step-by-step solvers can provide personalized help for students working through problems [ 28 ] , and aid educators in curriculum and course design by semantically relating exercises based on their solutions [ 22 , 23 ] . Indeed , studies have found automated tutors capable of yielding similar [ 3 , 20 ] or larger [ 29 ] educational gains than human tutors . While solving problems alone does not necessarily translate to good teaching , automated tutors typically have powerful domain models as their underlying foundation . However , even modest mathematical domains are challenging to solve . As an example , consider solving linear equations step-by-step using low-level axioms , such as associativity , reflexivity and operations with constants . This formulation allows all solution strategies that humans employ to be 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . expressed as combinations of few simple rules , making it attractive for automated tutors [ 28 , 25 ] . But while formulating the domain is simple , obtaining a general solver is not . Naïve search is infeasible due to the combinatorial solution space . As an example , the search-based solver used in the recent Algebra Notepad tutor [ 25 ] is limited to solutions of up to 4 steps . An alternative is manually writing expert solver heuristics . Again , even for a domain such as high-school algebra , this route is difficult and error-prone . As we describe in Section 5.2 , we evaluated Google MathSteps , a library that backs educational applications with a step-by-step algebra solver , on the equations from the Cognitive Tutor Algebra [ 28 ] dataset . MathSteps only succeeded in 76 % of the test set , revealing several edge cases in its solution strategies . Thus , even very complex expert-written strategies may have surprising gaps . An alternative could be to learn solution strategies via Reinforcement Learning ( RL ) . We formulate symbolic reasoning as an RL problem of deterministic environments that execute domain rules and give a binary reward when a problem is solved . Since we aim for generality , we assume a domainagnostic interface with the environment : states and actions are given to agents as unstructured text . These domains have several idiosyncrasies that make them challenging for RL . First , trajectories are unbounded , since axioms might always be applicable and lead to new states ( e.g . adding a constant to both sides of an equation ) . Second , agents have no direct access to the underlying structure of the domain , only observing strings and sparse binary rewards . Finally , each problem only has one success state ( e.g . x = number , in equations ) . These properties rule out many popular algorithms for RL . For instance , Monte Carlo Tree Search ( MCTS , [ 7 ] ) uses random policy rollouts to train its value estimates . If the solution state is unique , such rollouts only find non-zero reward if they happen to find the complete solution . Thus , MCTS fails to guide search toward solutions [ 1 ] . Indeed , as we show in Section 5.2 , Deep Q-Learning , and other algorithms that are based on estimating expected rewards , perform poorly in these symbolic domains . To overcome these challenges , we propose a novel learning algorithm , Contrastive Policy Learning ( ConPoLe ) , which succeeds in symbolic environments . Our key insight is to directly learn a policy by attempting to capture the mutual information between current and future states that occur in successful trajectories . ConPoLe uses iterative deepening and beam search to find successful and failed trajectories during training . It then uses these positive and negative examples to optimize the InfoNCE loss [ 24 ] , which lower bounds the mutual information between the current state and successful successors . This provides a new connection between policy learning and unsupervised contrastive learning . Our main contributions in this paper are : • We introduce 5 environments for symbolic reasoning ( Fig . 1 ) drawn from skills listed in the Mathematics Common Core Curriculum ( Section 3 ) . We find that existing Reinforcement Learning algorithms fail to solve these domains . • We formulate policy learning in deterministic environments as contrastive learning , allowing us to sidestep value estimation ( Section 4 ) . The algorithm we introduce , ConPoLe , succeeds in all five Common Core environments , as well as in solving the Rubik ’ s Cube ( Section 5.2 ) . • We provide quantitative and qualitative evidence that the problem representations learned by ConPoLe reflect the equation-solving curriculum from the Khan Academy platform . This result suggests a number of applications of representation learning in education . 2 Related Work . Automated mathematical reasoning spans several research communities . Theorem-proving languages , such as Coq [ 5 ] or Lean [ 10 ] , enable the formalization of mathematical statements and can verify proofs , but they are limited in their ability to discover solutions automatically . A rich line of recent work has focused on learning to produce formal proofs using supervised [ 27 , 32 , 2 ] and reinforcement learning [ 4 , 18 , 26 ] . Similar to GPT-f [ 27 ] , our model only assumes unstructured strings as input ; however , we do not use a dataset of human solutions . To the best of our knowledge , our work is the first to consider learning formal reasoning directly from text input ( like GPT-f ) by purely interacting with an environment ( like rlCop [ 18 ] and ASTatic [ 32 ] ) . Existing RL algorithms for theorem proving make use of partial rewards to guide search . For example , rlCop [ 18 ] learns to prove in the Mizar Verifier using Monte Carlo Tree Search . In Mizar , applying a theorem decomposes the proof into sub-goals . Random policy rollouts often close a fraction of the sub-goals , providing signal to MCTS . In our environments , however , there is only one solution state for each problem , causing MCTS to degenerate to breadth-first search . This challenge is similar in other large search problems , such as the Rubik ’ s Cube . DeepCubeA [ 1 ] handles the sparsity problem in the Rubik ’ s Cube by generating examples in reverse , starting from the solution state . This uses the fact that moves in the Rubik ’ s Cube have inverses , and that the solution state is always known a priori . In contrast , the algorithm we propose in Section 4 , which has neither of these assumptions , is still able to learn an effective Rubik ’ s Cube solver ; we compare it to DeepCubeA in Section 5.2 . Intelligent Tutoring Systems [ 20 , 28 , 22 ] provide a key application for step-by-step solvers . Tutors for symbolic domains , such as algebra [ 28 , 25 ] or database querying [ 3 ] , can use a solver to help students even in the absence of expert human hints ( unlike example-tracing tutors , which simply follow extensive problem-specific annotations [ 30 ] ) . Currently , solvers used in automated tutors are either search-based [ 25 ] , but limited to short solutions , or hand-engineered [ 17 ] . Our method can learn high-level solution strategies by composing simple axioms , generalizing to long solutions without the need for expert heuristics . Moreover , our analysis in Section 4 suggests that learned problem embeddings can be applied in mathematics education – another advantage of a neural solver . Finally , our approach to Reinforcement Learning in symbolic domains builds on unsupervised contrastive learning , specifically InfoNCE [ 24 ] . Contrastive learning has been used to learn representations that are then used by existing Deep RL algorithms [ 24 , 19 , 13 ] . We instead make a novel connection , casting policy learning as contrasting examples of successful and failed trajectories . 3 Setup . Motivated by mathematics education , we are interested in formal domains where problems are solved in a step-by-step fashion , such as equations or fraction simplification . While various algorithms could compute the final solution to such problems , finding a step-by-step solution by composing low-level axioms is a planning problem . Formally , we define a domain D = ( S , A , T , R ) as a deterministic Markov Decision Process ( MDP ) , where S is a set of states , A ( s ) is the set of actions available at state s , T ( st , a ) is the transition function , deterministically mapping state st ∈ S and action a ∈ A ( st ) to the next state st+1 , and R ( s ) is the binary reward of a state : R ( s ) = 1 if s is a “ solved ” state , otherwise R ( s ) = 0 . States with positive reward are terminal . Initial states can be sampled from a probability distribution PI ( s ) , which we assume to have states of varying distance to the solution state . This implicit curriculum allows agents to find a few solutions from the beginning of training – an important starting signal , since we assume no partial rewards . States in symbolic domains typically have a well-defined underlying structure . For example , mathematical expressions , such as equations , have a recursive tree structure , as do SQL queries and programming languages . A different , yet similarly well-defined structure dictates valid states in the Rubik ’ s cube . However , in order to study the general problem of symbolic learning , we want to assume no structure for states beyond the MDP specification . Therefore , we assume all states and actions are strings , i.e . S , A ( s ) ⊆ Σ∗ for some alphabet Σ . Naturally , our goal is to learn a policy π ( a|s ) that maximizes the probability that following π when starting at a randomly drawn state s0 ∼ PI ( · ) leads to a solved state sk .

Summary: This paper proposes a new reinforcement learning method for symbolic domains with sparse binary rewards (e.g., finding step-by-step solutions to math problems). The method assumes the availability of many problem instances within the domain, of varying levels of difficulty; the agent learns by solving the easier problems first, and training to distinguish positive and negative examples of state transitions on the way to a solution, using a version of the InfoNCE loss from the contrastive learning literature. ConPoLe operates directly on text representations of states within a domain, so adapting it to new domains does not require domain-specific hand-engineered features. The paper presents four novel domains corresponding to different kinds of math problem, inspired by the Common Core curriculum, and shows that ConPoLe outperforms existing methods by a wide margin. In addition, a small-scale experiment suggests that the embeddings that ConPoLe learns cluster according to existing pedagogical classifications of math problems on Kahn Academy. ConPoLe is also tested in an established Rubik's cube domain, comparing favorably to the DeepCubeA baseline in terms of data efficiency and search time (if not final solution length).