Source-linked AI summary
Reinforcement Learning for Symbolic Equation Solving
Kevin P O Keeffe
TL;DR
The paper asks whether reinforcement learning can produce replayable, step-by-step symbolic equation-solving procedures beyond closed equations. It formulates algebra as an MDP with dynamic actions and TreeMLP, and reports strong performance on closed equations plus controlled restricted-open families, while limiting open-equation claims to four benchmark families.
Problem
Most machine-learning symbolic-mathematics systems produce one-shot answers, leaving a need for explicit, replayable solution procedures.
Method
The paper casts algebra as an MDP with dynamic operation–sub-expression actions and a TreeMLP policy, training the main policy from reward alone while using a supervised, CAS-replaceable CoV generator.
Results
0.93 greedy accuracy matches ConPoLe’s 0.925 on CommonCore, while the restricted-open benchmark reaches 0.79 beam and 0.67 greedy accuracy.
Takeaways & Limitations
Reward-driven learning recovers a useful CoV trigger without a detector, with nontrivial timing content only for the nested-CoV exponential family.
Abstract
from arXiv · showhide
We present a reinforcement-learning agent that solves symbolic equations step by step, covering both nonlinear closed equations (radicals, exponentials, trigonometric) and a controlled class of restricted-open families requiring a change of variables (CoV) such as completing the square. We cast algebra as an MDP with a dynamic action space and a tree-structured policy (TreeMLP). The main policy learns from reward alone with no supervised solution traces; the CoV substitution comes from a supervised generator interchangeable with a CAS call. On closed equations the agent matches the prior best on CommonCore (0.93 greedy vs. ConPoLe's 0.925) under a single policy. On four hand-designed restricted-open families (quadratic, cubic, quartic, exponential) it reaches 0.79 beam / 0.67 greedy, exceeding the strongest non-learned search (A-star, 0.64). Learned CoV timing has content only on the exponential family, the one requiring a nested CoV, where a natural rule solves none of the held-out equations while the policy solves 75% from reward alone. At 10x scale a sharp seed-level bimodality emerges; a UCB learning-progress curriculum shows a non-significant positive trend toward mitigating it. We do not claim general open-equation solving: every open-equation result is confined to these four controlled families.
1 Introduction
The paper develops reward-driven, step-by-step algebraic equation solving, extending prior closed-equation work to a controlled restricted-open setting with change-of-variables operations.
- Motivation: Most machine-learning symbolic mathematics systems emit one-shot answers, whereas this work targets explicit, replayable transformation traces.Such traces support differential-equation reductions, theorem-prover tactics, and didactic solvers.
- Approach: The main policy learns equation-solving procedures from reward alone without supervised solution traces.The CoV substitution generator is the sole supervised component and can be replaced by a CAS call.
- Approach: The MDP uses state-dependent legal actions pairing algebraic operations with current equation sub-expressions, implemented with TreeMLP.TreeMLP supplies the policy/value network, while curriculum, success replay, and constant relabeling provide additional inductive biases.
- Results: 0.93 greedy accuracy on CommonCore matches ConPoLe’s 0.925 under a single policy.This result concerns the closed-equation benchmark.
- Results: Learned CoV timing matters only for the nested-CoV exponential family, where the policy solves 56/75 held-out equations versus 0/75 for a natural rule.On the three single-CoV families, a one-line hand-coded detector matches or beats the policy.
2 Related work
The paper distinguishes procedural, reward-driven equation manipulation from one-shot symbolic generation, formula synthesis, and proof-oriented reasoning, while adapting graph encoders to expression trees.
- RL for symbolic equation solving: Prior RL work addressed linear or nonlinear closed equations, while this paper adds restricted-open equations requiring change-of-variables structure.ConPoLe reports 92.5% on CommonCore, the benchmark used here.
- Neural symbolic math: Neural symbolic-math systems have used supervised one-shot generation, whereas this work learns from per-step rewards without supervised solution traces.The CoV policy component is supervised on known substitutions.
- Symbolic regression and program synthesis: Unlike symbolic regression and program synthesis, the agent manipulates a given equation through legal transformations.Its CoV macroaction is a concrete substitution rather than an induced abstraction library.
- Graph encoders for symbolic state: TreeMLP is designed for sparse, rooted expression trees rather than the large undirected neighborhoods assumed by GCN and GraphSAGE.It uses strictly local sum aggregation without degree normalization or learned edge transforms.
- Formal mathematical reasoning: The equation-solving setting is computational and complementary to theorem-proving systems that target proof-style validity.The related systems cited include Peano and AlphaZero-style theorem provers.
- Curriculum learning and capability collapse: The UCB-LP curriculum combines learning-progress sampling with a UCB1 bonus and addresses seed-level bimodality observed in symbolic reinforcement learning.The paper situates this failure mode alongside prior reports of bimodal capability acquisition.
3 Problem Formulation
The formulation represents symbolic equation solving as an MDP with dynamic operation–sub-expression actions, TreeMLP state encoding, and terminal verification of one mapped-back root.
- 3 Problem Formulation: The MDP state is an equation, and actions pair an algebraic operation with a current sub-expression.The action list is re-derived dynamically from the equation state.
- 3 Problem Formulation: Sub-expressions from both equation sides form the variable-length term set, which supports operations such as subtraction, division, and expansion.The action space is capped at |A| = 50, with illegal actions masked at each step.
- 3 Problem Formulation: Reward is based on reducing equation complexity measured by expression-tree nodes and edges.For example, C(ax + b) = 9 and C(ax^2 + bx + c) = 19.
- 3 Problem Formulation: The environment applies exact SymPy transitions to tracked left- and right-hand sides and terminates when x equals a substitution-verified right-hand side.Stored inverse transformations map CoV or relabeled answers back to the original variables.
- 3 Problem Formulation: Success means producing one terminally verified root, not recovering the complete solution set.The tracked left side must be x, the right side must contain no x, and substitution must simplify the active equation to 0.
- 3 Problem Formulation: Intermediate operations are replayable and inspectable but are not individually checked for equivalence, so traces are not proof objects.SymPy exceptions and canonicalizer operation limits count as failures.
- 3 Problem Formulation: Restricted-open equations require adding out-of-term-set structure, such as completing the square, and are distinct from closed equations.The benchmark studies only controlled families admitting a known depression substitution.
- 3.1 TreeMLP: a tree-structured policy network: TreeMLP embeds node IDs, performs local sum-aggregation message passing, and pools valid nodes for PPO policy and value heads.Its design avoids degree normalization and learned edge transforms for sparse expression trees.
4 Results
The agent extends step-by-step reinforcement learning to controlled restricted-open equations while retaining strong closed-equation performance. Results show policy-dominated open-equation accuracy, with learned CoV timing contributing specifically on nested exponential problems.
- Closed equations: 0.93 greedy test accuracy on CommonCore matches ConPoLe’s 0.925 under the closed-equation inductive-bias setup.TreeMLP and success-replay contribute to performance and scaling across the closed-equation benchmarks.
- Scope: The open-equation claims are limited to four hand-designed restricted-open families admitting known depression substitutions, not general open-equation solving.Generic same-degree polynomials without such a fixed CoV are excluded by construction.
- Restricted-open results: 0.79 testbeam and 0.67 greedy accuracy are achieved on the headline open_small cohort across four restricted-open families.The headline cohort uses five validation-selected seeds and separates these figures from cohort W’s 0.73 greedy result.
- Policy versus search: 0.67 greedy accuracy exceeds A∗ up to 0.64 and BFS 0.26, while beam decoding adds approximately 0.12 to reach 0.79.The comparison uses the same policy and action space on cohort H; the result is policy-dominated rather than search-dominated.
- CoV trigger: The exponential family is the only setting where learned CoV timing has content: the policy solves 56/75 cases, whereas the natural rule solves 0/75.The family requires two CoV steps; quadratic, cubic, and quartic use a single-CoV recipe and are solvable with an immediate detector.
- CoV trigger: The learned trigger recovers a hand-engineered trigger narrowly, with aggregate trigger accuracy between the naive and expand-preconditioned rules.The paper frames this as recovery rather than superiority, because the stronger rule uses a precondition discovered by inspecting learned traces.
5 Conclusion
The paper extends reinforcement learning for step-by-step symbolic equation solving to a controlled restricted-open setting, with the main policy trained from reward alone. Learned CoV timing matters specifically for the nested-CoV exponential family.
- The agent solves controlled restricted-open equations step by step with a reward-trained main policy and a supervised, CAS-replaceable CoV generator.
- On the nested-CoV exponential family, the natural rule solves 0/75 while the policy solves 56/75 from reward alone.
- On the three single-CoV families, a one-line rule matches or beats the policy, so learned timing has content only for the exponential family.
6 Limitations
The paper’s limitations concern evaluation uncertainty, restricted solvability, reproducibility, and the indirect value of solving equations whose underlying solutions are CAS-accessible. Proposed next steps include retraining, replicated evaluation, curriculum improvements, and integrating search into training.
- Evaluation uncertainty: α = 0 removes a penalty in the open_small control, raising full-stack performance from 0.490 to 0.891 with 0/9 stalls.The authors did not promote this setting because downstream analyses used α = 0.1 checkpoints.
- Evaluation uncertainty: 0.79 is not claimed as a lower bound because the seed-matched protocol produced 0.490 ± 0.377 with a 1/3 stall rate.Retraining at α = 0 under the headline protocol is identified as the next step.
- Evaluation uncertainty: Seed matching removes seed-selection confounding but not residual run-to-run variance, measured at 0.079 on average and 0.253 in the worst case.Without per-cell replicates, the ablation resolves only large sign-stable effects.
- Scope and reproducibility: The equations are solvable by SymPy in milliseconds, so the contribution is the RL formulation and replayable step-by-step traces rather than new bare equation-solving capability.The paper positions these traces for tutoring, differential-equation solving, and theorem-prover tactics, while not releasing the agent or trained checkpoints.
- Learning stability: UCB-LP shows a non-significant positive trend, with 6/8 escape and two seeds still stalling; open_large retains a 2/8 stall rate and a 0.09 gap to oracle.The authors identify tighter curricula and a hardened numerical pipeline as next steps.
- Structural scope: The action set bounds the solvable class: the cubic family required a cube-root action, and analogous gaps may occur beyond the four CoV templates.The abel_level3 class is weakest, with 0.53 full-stack performance on open_small and an escape-seed mean of 0.39 on open_large.
- Evaluation and training: 0.67 → 0.79 on cohort H and 0.73 → 0.84 on cohort W shows a residual decode-time search gap because search is used only during evaluation.Folding search into training is proposed as a promising direction.
A.1 Macroactions
Three macroactions are required for the considered symbolic transformations: expansion, collection of a common variable, and multiplication by −1. The examples show that these actions enable distribution, factoring, and isolation steps unavailable from the basic action set.
- Required transformations: The (expand, None) macroaction distributes a coefficient across a bracketed linear term.Without it, dx + c(ax + b) = e cannot be transformed into dx + cax + cb = e.
- Required transformations: The (collect, x) macroaction factors the common x from the first two terms after expansion.It is used on dx + cax + cb = e.
- Required transformations: The (mul, −1) macroaction isolates x in −x = b because −1 is absent from the symbolic term set.The authors retain a purely symbolic term set instead of adding −1 directly.
A.2 Dataset generation: rational equations
The authors supplement recursive datasets with non-degenerate rational equations using symbolic coefficients, while limiting their proportion to preserve diversity across functional forms. Several retained forms have explicitly listed closed-form solutions.
- Construction and filtering: The rational-equation dataset uses forms with non-zero symbolic coefficients a, b, and c.Degenerate cases, including identically vanishing denominators or trivial numerator cancellation, are excluded.
- Retained forms: Examples include rational equations whose solutions are x = −b, x = −b − ab, and x = −b + ac.These are presented as retained equation forms.
- Dataset mixture: The rational forms are merged into both small and large datasets with limited proportion to preserve diversity across logarithmic, trigonometric, exponential, and other forms.The limitation applies to the mixture rather than excluding rational equations entirely.
A.3 Illegal actions
The solver prevents both explicit and hidden divisions by zero by removing unsafe division actions from the dynamic action set.
- The action filter removes division by zero and divisions by polynomial factors that may vanish at solutions.This handles hidden zero denominators such as dividing by x+a when (x+a)(x+b)=0.
A.4 Hyperparameters and compute
The experiments use Stable Baselines 3 PPO defaults and specified architecture constants, running entirely on CPU with substantial per-run compute for open_large.
- Stable Baselines 3 default hyperparameters are used for PPO alongside constants specified in the body and configuration tables.
- All experiments run on a single 24-core CPU machine without a GPU.The PPO device setting resolves from auto to CPU.
- Each open_large training uses 5×10^6 steps, while open_small headline and seed-matched controls use 3×10^6 steps.The beam-width and trigger cohort uses 5×10^6 steps.
- The action-diversity penalty is α=0.1 for open_small headline cohorts, while open_large, curriculum, intervention, and closed CommonCore runs use α=0.Separate anti-loop sweep rows deliberately vary α.
A.5 Closed-equation analysis
Closed-equation performance is strong on shallow small data but degrades substantially with deeper nesting, with failures concentrated in transcendental and deep polynomial cases.
- Per-type accuracy: ≥0.88 accuracy is achieved on linear and radical equations, while degree-2–4 polynomial classes reach 0.75–0.82 on the small dataset.The trigonometric bucket reaches 72% accuracy and accounts for most remaining failures.
- Per-type accuracy: The small-dataset polynomial buckets contain only depressed forms; open polynomials with linear cross-terms are evaluated separately.
- Failure modes: 75 of 387 failed cases involve trig inverse-function steps, transcendental simplification, or exceeding the action budget on deep polynomials.
- Scaling: Accuracy falls across every category on the large dataset, including poly-deg4 from 0.75 to 0.08.The larger set adds substantially deeper nesting rather than merely more examples.
- CommonCore: 0.93 greedy test accuracy on CommonCore matches ConPoLe’s 0.925 under matched greedy decoding.Restricting to non-degenerate templates raises the non-buffer variant to 0.955.
A.6 Open-equation embeddings (converged vs. stalled seeds)
The open-equation traces show a shared change-of-variables recipe, with nested substitution unique to exponential equations, while broader probes relate performance to policy behavior rather than encoder collapse.
- CoV generation: The CoV generator decodes substitutions as self-delimiting prefix grammar productions, with production arities determining structure.
- A.8 Open-equation solution traces: The four change-of-variables classes share a seven-step recipe up to constants, while exponential equations alone use a nested CoV.The closed-form abel_level3 example uses no CoV and solves in three steps.
- Training design: Sparse outcome-only rewards substantially degrade performance and require a one-second-per-step timeout to keep training tractable.
- Encoder comparison: TreeMLP outperforms GCN and GraphSAGE on the small dataset, although the single-seed comparison is suggestive rather than conclusive.
- Ablations: Curiosity bonuses do not improve TreeMLP learning on abel_level3, based on single-seed evidence.
- Stalled seeds: The stalled-seed probes find no evidence of capacity loss, data-distribution repair, or encoder representation collapse.
- Stalled seeds: The evidence is most consistent with a policy-attractor phenomenon, while not exhaustively ruling out other explanations.