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: In this paper, the authors address the problem of learning policies on symbolic domains - like mathematical equations. They propose ConPoLe - learning algorithm that relies on optimizing InfoNCE, a well-known contrastive loss. They benchmark ConPoLe on four different environments and show that it generalizes well. A key contribution of the paper is the introduction of these environments themselves.