Interactive Symbolic Regression through Offline Reinforcement Learning: A Co-Design Framework

TL;DR

This work tackles the SR challenge posed by vast expression spaces by introducing Sym-Q, an offline RL framework that builds symbolic expressions via a tree-structured action space and a modular encoder architecture. A key innovation is the co-design mechanism, which allows domain experts to iteratively inject priors and modify expression components to guide discovery toward physically meaningful models. Sym-Q combines a point-set encoder, a tree encoder, a Q-network, and a Conservative Q-learning objective, augmented with supervised contrastive learning, enabling robust skeleton recovery and high fitting accuracy across benchmarks, notably SSDNC where skeleton recovery reaches 82.3% with beam search and $R^2$ around $0.951$. The approach avoids transformer-decoder dependencies, offers flexible encoder choices, and demonstrates strong performance gains from interactive domain knowledge, with promising implications for real-world scientific discovery and streamlined model discovery in physics-informed contexts.

Abstract

Symbolic Regression (SR) holds great potential for uncovering underlying mathematical and physical relationships from observed data. However, the vast combinatorial space of possible expressions poses significant challenges for both online search methods and pre-trained transformer models. Additionally, current state-of-the-art approaches typically do not consider the integration of domain experts' prior knowledge and do not support iterative interactions with the model during the equation discovery process. To address these challenges, we propose the Symbolic Q-network (Sym-Q), an advanced interactive framework for large-scale symbolic regression. Unlike previous large-scale transformer-based SR approaches, Sym-Q leverages reinforcement learning without relying on a transformer-based decoder. This formulation allows the agent to learn through offline reinforcement learning using any type of tree encoder, enabling more efficient training and inference. Furthermore, we propose a co-design mechanism, where the reinforcement learning-based Sym-Q facilitates effective interaction with domain experts at any stage of the equation discovery process. Users can dynamically modify generated nodes of the expression, collaborating with the agent to tailor the mathematical expression to best fit the problem and align with the assumed physical laws, particularly when there is prior partial knowledge of the expected behavior. Our experiments demonstrate that the pre-trained Sym-Q surpasses existing SR algorithms on the challenging SSDNC benchmark. Moreover, we experimentally show on real-world cases that its performance can be further enhanced by the interactive co-design mechanism, with Sym-Q achieving greater performance gains than other state-of-the-art models. Our reproducible code is available at https://github.com/EPFL-IMOS/Sym-Q.

Figures (12)

• Figure 1: Overview of the proposed framework. a). The expression and its corresponding expression tree. b). The proposed Sym-Q agent supports both offline training with ground truth human knowledge and potentially online searching with reward signals. $\tau_{circle}$ represents the agent trajectories within the symbolic environment. c). The Sym-Q architecture and step-wise decision-making mechanism.
• Figure 2: The figures illustrate enhanced model performance on the SSDNC dataset when partial ground truth of the equation is provided by domain experts ($R^2$ (a, b) and recovery rate (c, d) ). The x-axis represents the percentage of given action sequences relative to the total. We analyzed equations of varying lengths, excluding categories with fewer than five samples to ensure statistical validity. Notably, both average $R^2$ and recovery rates increase as more ground truth sequence steps are incorporated, demonstrating that domain expert guidance significantly improves model performance.
• Figure 3: Overall performance improvement with respect to the ratio of ground truth information provided for Sym-Q and NeSymReS, showing the co-design impact with up to a maximum of 60% additional information provided.
• Figure 4: Model accuracy in retrieving different additive and multiplicative drifts over the AI Feynman benchmark. Violet (NeSymRes) and magenta (Sym-Q) bars depict the recovery rate of equations with additional drifts, without giving any prior knowledge. Salmon (NeSymRes) and yellow (Sym-Q) bars refer to the case where the textbook equation is given as a prior and, therefore, only the additional components need to be recovered.
• Figure 5: Correlation between the transit spectra dataset marquez2018supervised entries $\pi_1$, $\pi_2$ and the observation values $f$ and $f\pi_1$. These were used to set up the two priors we used to co-design the analytical expression of $f$.