[Paper Review] Contrastive Reinforcement Learning of Symbolic Reasoning Domains

Gabriel Poesia, WenXin Dong+1 more|arXiv (Cornell University)|Jun 16, 2021

Intelligent Tutoring Systems and Adaptive Learning32 references4 citations

TL;DR

This paper introduces Contrastive Policy Learning (ConPoLe), a novel reinforcement learning algorithm that improves symbolic reasoning in domains like mathematics by optimizing an InfoNCE contrastive loss to maximize mutual information between current and future states on successful solution paths. ConPoLe achieves strong performance across four Common Core-inspired math environments and the Rubik’s Cube, learning semantically meaningful problem representations that predict curriculum categories with 90.5% accuracy.

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

Motivation & Objective

• To address the challenge of training reinforcement learning agents in symbolic reasoning domains—such as algebra and logic—where sparse binary rewards and unstructured text states make learning difficult.
• To overcome limitations of existing RL algorithms, which fail in these domains due to unbounded trajectories, lack of structural insight, and sparse rewards.
• To develop a domain-agnostic, text-based RL framework that enables generalization across symbolic domains without handcrafted features or human solutions.
• To learn problem representations that reflect semantic curriculum structure, enabling accurate categorization of math problems.
• To demonstrate the generality and effectiveness of the approach on both educational math problems and complex search problems like the Rubik’s Cube.

Proposed method

• Formulate symbolic reasoning domains as deterministic, text-based RL environments where states and actions are unstructured strings and only binary rewards indicate solution success.
• Use iterative deepening and beam search during training to collect positive (successful) and negative (failed) trajectories for contrastive learning.
• Optimize the InfoNCE loss to maximize mutual information between current states and their successful future states, effectively learning a policy without value function estimation.
• Leverage contrastive learning to learn meaningful state representations that capture semantic structure in problem-solving paths.
• Apply the trained policy with Batch Weighted A* Search (BWAS) for inference, using predicted log-probabilities as node weights in search.
• Train a single model architecture across diverse domains, including math equations and the Rubik’s Cube, using a unified text-to-text interface.

Experimental results

Research questions

• RQ1 Can contrastive learning be effectively applied to policy learning in symbolic reasoning domains with only binary rewards and unstructured text?
• RQ2 Can a contrastive approach outperform standard RL baselines in symbolic domains where value estimation fails due to sparse rewards and long trajectories?
• RQ3 Do the representations learned by the agent reflect semantic curriculum structure, such as problem categories in a real educational platform?
• RQ4 Can the same algorithm generalize across diverse symbolic domains, including math equations and the Rubik’s Cube, without architectural or reward function changes?
• RQ5 Is it possible to learn effective solvers for symbolic domains without access to human solutions or hand-engineered features?

Key findings

• ConPoLe successfully solves all four Common Core-inspired symbolic reasoning environments, whereas standard RL baselines fail.
• The problem representations learned by ConPoLe achieve 90.5% accuracy in predicting the category of math problems from the Khan Academy curriculum, significantly outperforming baseline methods and even surpassing string edit distance.
• ConPoLe’s learned representations form highly structured clusters corresponding to sections of the Khan Academy curriculum, despite no explicit curriculum supervision.
• On the Rubik’s Cube, ConPoLe solves all 100 test instances scrambled with 1000 moves, matching DeepCubeA’s success rate but visiting 3 million nodes on average—36% fewer than DeepCubeA’s 8.3 million.
• ConPoLe’s solutions are longer (39.4 moves on average) than DeepCubeA’s (21.6 moves), but this difference is expected given ConPoLe was trained on 100x fewer environment steps than DeepCubeA.
• The success of ConPoLe demonstrates that contrastive learning can effectively replace value function estimation in symbolic planning, enabling sample-efficient policy learning in sparse-reward, unstructured environments.