Curiosity-driven RL for symbolic equation solving

AI summary

Overview

Research area: Reinforcement learning applied to symbolic mathematics (computer algebra / equation solving).

Technical level: Advanced. The reader needs familiarity with RL algorithms (PPO, A2C), curiosity-driven exploration (ICM, RIDE, NGU, RND), and symbolic expression trees.

Scope: This paper tests whether model-free PPO, augmented with curiosity-based exploration and a graph-based (expression-tree) action space, can learn to solve nonlinear single-variable algebraic equations — radicals, exponentials, logarithms, and trigonometric functions — that go beyond the linear equations studied in prior work.

What This Paper Is About

Symbolic mathematics is still dominated by hand-engineered systems like Mathematica or Maple, and applying RL to it is hard because the state space is combinatorially large (equations branch into many sub-expressions) and the action space is large and dynamic. The paper asks whether RL agents can learn algebraic manipulation autonomously, and shows that curiosity-driven exploration plus an expression-tree action representation lets agents solve nonlinear equations — including ones "closed" enough that no new terms must be invented.

Key Contributions

1. A graph-based MDP formulation for symbolic equations. Equations are represented as expression trees vectorized by preorder traversal (padded to length L = 50), and actions are high-level (operation, term) pairs drawn from a dynamic list of sub-expressions, rather than the primitive, character-level axioms used by Poesia et al. (2021).

2. Three macroactions shown to be necessary, not merely convenient. (Expand, None), (collect, x), and (multiply, -1) are added to the action set; the appendix argues each is required to solve the equations considered (for example, distributing c into ax+b in dx+c(ax+b)=e, factoring a common x, and handling -x=b when -1 is not in the term list).

3. A dense complexity-based reward. Reward is defined as the drop in equation complexity C, counting nodes plus edges in the expression tree, so C(ax+b) = 5+4 = 9 and C(ax²+bx+c) = 10+9 = 19. This replaces the sparse binary solved/unsolved signal of prior work.

4. An empirical comparison of four curiosity methods on symbolic equations. PPO was combined with ICM, RIDE, NGU, and RND, and benchmarked against A2C, plain PPO, and a non-learning A* baseline on both fixed and random equation environments.

Main Findings

• Curiosity-based exploration is required for hard equations. In the fixed-equation setting, all methods solved ax+b and a/x+b at 1.00 success, but performance diverged on nested equations. On the most challenging two, (ax+b)+e(cx+d) and e+(ax+b)/(cx+d), PPO-RND achieved 0.40 and 0.50 while A2C, PPO, and PPO-RIDE scored 0.00 on both.

• RND was the best-performing curiosity method overall. On c+d/(ax+b), PPO-RND led with 0.80, followed by PPO-ICM and PPO-NGU at 0.50, while the A* baseline failed at 0.00. On c(ax+b)+d, PPO-RND reached 1.00 compared with PPO-ICM 0.70, PPO-RIDE 0.60, and PPO-NGU 0.40.

• PPO-NGU and PPO-RND were the top two performers in the fixed environment, while A* solved the two easiest equations (1.00) and failed on several harder ones.

• Generalization on random equations depends heavily on the agent. On both the small (depth < 4, size 3874, sub-sampled to a 10³/10² train/test split) and large (depth < 5, size 15625, sub-sampled to a 10⁴/10³ train/test split) datasets, PPO-RND and PPO-RIDE were the best performers, reaching test₁₀ ≈ 1.0 and 0.8 respectively.

• Greedy performance lagged stochastic performance. Test₁₀ success was much higher than test_greedy ≈ 0.25, meaning the agents often needed multiple stochastic rollouts to find a solution.

• PPO failed to scale, and PPO-NGU regressed. Plain PPO showed decent performance on the small dataset but failed on the large one entirely; PPO-NGU was the worst performer in the random-equation environment, showing little progress on either dataset. The discussion attributes NGU's weakness to its higher memory footprint.

• The MDP only handles "closed" equations. Solving requires manipulating terms already present in the equation. "Open" equations — such as the quadratic, solved by completing the square with (b/2a)², a term not in the term set — are outside the scope of this work, as they were in prior work by Poesia et al. (2021) and Dabelow & Ueda (2024).

Methodology in Plain English

The environment is a single-variable algebra problem. An equation is turned into a tree of symbols, flattened into a fixed-length list of integers, and split into left-hand and right-hand sides. The agent's legal moves are pairs like (subtract, b) or (divide, a), where the second element must be one of the sub-expressions already sitting in the equation. Three special macro-moves (expand a product, collect like powers of x, multiply both sides by −1) were added and justified as necessary. SymPy applies each chosen operation, and the agent receives a reward equal to how much the equation's node-and-edge count dropped. The episode ends when the left side is x and substituting the right side back into the original equation simplifies to 0. Illegal moves such as division by zero — including "hidden" division by a factor that could be zero — are masked out. The action set was capped at 50 entries; the authors also tried 40, 70, and 100 and report no change in performance.

Training used A2C and PPO from Stable Baselines 3 with default hyperparameters, and curiosity modules (ICM, RIDE, NGU, RND) from the rllte package. Each algorithm was trained for N_train = 3×10⁶ steps over N_trial = 10 trials, scored by mean success rate ± standard error. A homemade A* search served as a non-learning baseline, using the complexity measure as a heuristic with depth limited to 20. For the random-equation experiments, datasets were built by starting with x and recursively applying actions, keeping terms restricted to (a, b, c), discarding anything SymPy could not solve, and counting multi-root equations as solved when any single root was found. Because that generation process does not produce rational equations like ax+b/(cx+d)+e=0, the authors augmented both datasets with a subset of such equations.

Why This Matters

The paper's claim is that RL can move beyond the linear equations of prior work and solve nonlinear equations involving radicals, logarithms, exponentials, and trigonometric functions — something the authors state has never been done before with RL. It also argues that solving these equations is a key building block for more advanced tasks like solving differential equations, and that an RL approach could reduce manual curation in symbolic mathematics and potentially find novel solution techniques.

Real-world applications (the paper does not enumerate commercial applications; these follow from the mathematics it targets):

• Computer algebra systems such as Mathematica or Maple, which currently rely on hand-engineered rules the authors say an RL agent could partly learn autonomously.
• Automated equation solving inside tutoring or homework-checking tools, since the fixed-equation test cases are drawn from elementary algebra textbooks.
• Scientific and engineering workflows where nonlinear single-variable equations must be rearranged symbolically before numerical solution.
• Downstream symbolic tasks such as analytic differentiation and integration and differential equation solving, which the paper identifies as the next building block.

Industry relevance: The strongest signal is methodological. The paper reports that RND is the best and lightest of the curiosity methods tested, while NGU's higher memory footprint may explain its poor performance, which is directly useful to practitioners choosing exploration modules. The finding that the action-set size could be varied between 40 and 100 with no performance change also suggests the approach is not fragile to that hyperparameter.

Future Directions

• Handle "open" equations. Extending the method to equations requiring new terms, auxiliary expressions, or changes of variables will likely require a generative sub-model, since the current action tuple always takes its term from the existing expression tree.
• Tackle the polynomial benchmarks in Box 1. The quadratic ax²+bx+c=0 ought to be within reach, while the cubic and quartic require discovering Cardano's and Ferrari's methods. The authors speculate that an AlphaZero-like algorithm, which excels at long-horizon planning, might help.
• Investigate NGU's weakness. The higher memory footprint of NGU may explain its poor performance on the random equation environment, and the authors suggest future work could explore this issue.
• Improve greedy vs. stochastic gap. Since test₁₀ ≈ 1.0 and 0.8 but test_greedy ≈ 0.25, closing the gap between stochastic rollouts and greedy deployment is an open practical question.

Target Audience

RL researchers interested in exploration methods and non-standard action spaces; symbolic computation and automated reasoning researchers curious about learned alternatives to hand-engineered algebra rules; and graduate students or practitioners working at the intersection of machine learning and mathematics. Readers without RL background can follow the problem framing and results tables, but the methodology section assumes working knowledge of PPO, A2C, and intrinsic-reward exploration.

Authors’ abstract

We explore if RL can be useful for symbolic mathematics. Previous work showed contrastive learning can solve linear equations in one variable. We show model-free PPO \cite{schulman2017proximal} augmented with curiosity-based exploration and graph-based actions can solve nonlinear equations such as those involving radicals, exponentials, and trig functions. Our work suggests curiosity-based exploration may be useful for general symbolic reasoning tasks.