Bayesian Symbolic Regression with Entropic Reinforcement Learning

Description

This is the official repository to the paper "Bayesian Symbolic Regression with Entropic Reinforcement Learning" by Oussama Boussif, Mohammed Mahfoud, Younesse Kaddar, Moksh Jain, Sida Li, Damiano Fornasiere, Xiaoyin Chen, Yoshua Bengio and Esmeralda S. Whitammer.

ERRLESS is a symbolic regression framework that trains a policy to sample mathematical expressions from a posterior over symbolic equations, rather than searching for a single best fit. Expressions are generated autoregressively as post-order traversals of expression trees by a transformer policy, with continuous constants fit jointly with the discrete skeleton. Sampling proportionally to a reward built from data likelihood and a structural yields not just the best equation but a diverse set of plausible equations.

Installation

Install uv

This repository uses uv to manage dependencies — install it if you haven't already.

Create the virtual environment

After cloning the repository, create and activate an environment:

uv venv
source .venv/bin/activate

Install the packages

uv sync

Getting Started

Single run

Runs are launched through main.py and configured entirely via Hydra. The experiment/ configs are the primary entry points — they wire together an algo, an environment, and training settings:

uv run main.py experiment=tbgfn_cont environment=feynman_I_6_2 seed=0

Any config value can be overridden from the CLI, e.g.:

uv run main.py experiment=tbgfn_cont environment=feynman_I_6_2 \
  environment.noise_level=0.01 algo.forward_policy_lr=1e-3 trainer.max_steps=2000

Algorithms

Algorithm	Config key (experiment=)	Description
Trajectory Balance	tbgfn_cont	TB-GFN with continuous constant parameters
Detailed Balance	dbgfn_cont	DB-GFN with continuous constant parameters
PySIPS (SMC)	sbatch_scripts/pysips.sh / scripts/run_pysips.py	Sequential Monte Carlo baseline (standalone script, not a Hydra algo)

Datasets

Dataset configs live in configs/environment/ and are passed via environment=<name>. Available families:

• Feynman — feynman_I_*, feynman_II_*, feynman_III_*, feynman_test_* (AI Feynman equations)
• PMLB "blackbox" — e.g. 579_fri_c0_250_5, 1028_SWD, 1089_USCrime, 1193_BNG_lowbwt (real/semi-synthetic regression datasets from PMLB)
• Synthetic — synthetic_1 … synthetic_7 (generated ground-truth expressions)

Dataset generation is documented in the corresponding notebooks (notebooks/blackbox_datasets.ipynb, notebooks/synthetic_dataset.ipynb).

Priors

The reward combines a likelihood (gfn_sr/logdist/gaussian.py) with a structural prior over expression trees, selected via environment.logprior=<name>:

• unigram_prior — independent per-token frequency prior
• nodes_prior — soft penalty on expression size
• pcfg_prior — probabilistic context-free grammar prior fit on a corpus of expressions (notebooks/pcfg_fit.ipynb, weights in pcfg.pt)
• uniform_prior — flat prior over valid expressions

Experiments

SLURM launch scripts for reproducing benchmark runs live in sbatch_scripts/:

bash sbatch_scripts/feynman.sh <wandb_tag>
bash sbatch_scripts/strogatz.sh <wandb_tag>
bash sbatch_scripts/blackbox.sh <wandb_tag>
bash sbatch_scripts/synthetic.sh <wandb_tag>

Baselines:

bash sbatch_scripts/pysips.sh                        # PySIPS (SMC) baseline

Ablations:

bash sbatch_scripts/ablations/db_ablation.sh          # Detailed Balance vs Trajectory Balance
bash sbatch_scripts/ablations/prior_ablations.sh      # structural prior comparison
bash sbatch_scripts/ablations/rb_ablation.sh          # replay buffer ablation
bash sbatch_scripts/ablations/no_dimensional_analysis.sh
bash sbatch_scripts/ablations/no_op2op_constraints.sh
bash sbatch_scripts/ablations/soft_length_prior.sh

Hyperparameter tuning

We use Optuna (via hydra-optuna-sweeper) for HPO. The sweep config is at configs/hparams_search/sr_nodes_optuna.yaml.

Single worker

uv run main.py -m seed=0 hparams_search=sr_nodes_optuna \
  experiment=tbgfn_cont \
  environment=feynman_I_6_2 \
  environment.noise_level=0.01 \
  environment.max_nodes=30 \
  algo.n_trajectories=1 \
  algo.reward_temperature_scheduler=null \
  algo.forward_policy.max_seq_len=31 \
  algo.forward_policy.vocab_size=17 \
  algo.forward_policy.output_dim=16 \
  trainer.max_steps=625 \
  resume=False \
  environment.batch_size=800

Distributed (multiple SLURM workers, shared study)

# 8 workers, default dataset (feynman_I_6_2)
./sbatch_scripts/submit_distributed_tuning.sh 8

# 12 workers, custom dataset
./sbatch_scripts/submit_distributed_tuning.sh 12 feynman_I_8_14

Each dataset gets its own Optuna database (optuna_study_<dataset>.db) and log directory (logs/sr_distributed_tuning_<dataset>/). Initialize a fresh study with uv run python initialize_distributed_study.py before the first submission.

Configuration & paths

• PROJECT_ROOT — repository root; all config paths resolve relative to it (${paths.root_dir}, set via rootutils).
• SCRATCH — base directory for logs/outputs (configs/paths/default.yaml); falls back to the working directory if unset.
• resume: True by default in configs/main.yaml — reruns with the same output directory auto-resume from the last checkpoint.
• Logging defaults to Weights & Biases; override with logger=csv or logger=tensorboard.