ChatterjeeLab / DooABLe

DooABLe

Learning Target Distributions with Budgeted Executable Paths

DooABLe learns a specified distribution over reachable outcomes together with cost-dependent execution paths. The implementation supports finite action graphs, bounded string editing, lattice navigation, and reaction-based molecular construction. Molecular generation begins with a supplied parent inventory and permits a chosen number of additional reactions.

For outcome weights $r(y)>0$, path cost $C(\xi)$, and route temperature $\tau>0$,

$P^{\ast }\left(\xi \right)=\frac{r\left(y\right)}{\underset{y^{\prime }}{\sum }r\left(y^{\prime }\right)}\frac{e^{-C\left(\xi \right)/\tau }}{\underset{{\xi }^{\prime }:y\left({\xi }^{\prime }\right)=y}{\sum }{e}^{-C\left({\xi }^{\prime }\right)/\tau }}.P^*(\backslash xi)=\backslash frac\{r(y)\}\{\backslash sum\_\{y\text{'}\}r(y\text{'})\}\backslash frac\{e^\{-C(\backslash xi)/\backslash tau\}\}\{\backslash sum\_\{\backslash xi\text{'}:y(\backslash xi\text{'})=y\}e^\{-C(\backslash xi\text{'})/\backslash tau\}\}.$ P ∗ ( ξ ) = ∑ y ′ ​ r ( y ′ ) r ( y ) ​ ∑ ξ ′ : y ( ξ ′ ) = y ​ e − C ( ξ ′ ) / τ e − C ( ξ ) / τ ​ .

The first factor specifies outcome probability. The second distributes that probability among feasible routes to the outcome. Prefix free energies are fitted through a backward recursion, and the forward sampler is trained with trajectory balance.

Install

Python 3.12 or newer is required. The executed checks used Python 3.12 and CPU PyTorch. The small examples run on a CPU and fit comfortably in a few GB of RAM.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
python -m pip install -e '.[test]'

The exact versions used for validation are in docs/environment.json . To reproduce those versions, install requirements-validated.txt after the CPU PyTorch installation.

Shortest example

dooable toy --output results/demo/graph.json
dooable train --graph results/demo/graph.json --steps 1200 --output results/demo/model
dooable sample --model results/demo/model --n 1000 --output results/demo/samples.jsonl

This example has two equally weighted outcomes and unequal route counts. Training writes the graph, model weights, optimizer state, and progress measurements. Each sample contains its outcome, edge sequence, action records, total cost, and log probability.

The exact dynamic program uses the same graph format.

dooable exact --graph results/demo/graph.json --temperature 0.7 --output results/demo/exact
dooable sample --model results/demo/exact --n 1000 --output results/demo/exact_samples.jsonl

Public molecular workflow

The included inventory has 12 named parent compounds, 13 reagents, and four directional amide/ester templates. A budget of two produces 168 distinct outcomes under the declared reaction rules and 750 Da product cap. These compounds provide a reproducible chemistry example. Enamine experiments use the supplied local catalog through the same adapter.

dooable build --parents data/examples/parents.csv --reagents data/examples/reagents.csv --reactions data/examples/reactions.json --budget 2 --output results/public/graph.json

The public property workflow downloads Caco-2 Wang data through Therapeutics Data Commons and BACE inhibition data from MoleculeNet. It canonicalizes structures, averages duplicate labels, creates scaffold-disjoint partitions, and fits Morgan-fingerprint extra-trees predictors. The saved predictor weights are included at results/validation/property_models/ . Refit them with

dooable fit-properties --data data/downloads --output results/public/property_models --seed 0

Use the included weights for the following generation run.

dooable score --graph results/public/graph.json --models results/validation/property_models --output results/public/rewards.json
dooable train --graph results/public/graph.json --rewards results/public/rewards.json --temperature 0.7 --steps 1200 --output results/public/model
dooable sample --model results/public/model --n 1000 --output results/public/samples.jsonl
dooable replay --samples results/public/samples.jsonl --budget 2
python scripts/evaluate_candidates.py --samples results/public/samples.jsonl --scores results/public/rewards.csv --budget 2 --requested 1000 --oracle-calls 168 --method dooable --output results/public/evaluation.json

BACE scores are predicted pIC50 values. Caco-2 scores are predicted log10 permeability in cm/s. Reaction replay checks the recorded graph transformations. Laboratory reaction success and experimental property measurements require separate experiments.

Benchmarks and plots

python scripts/run_benchmarks.py --config configs/toy.yaml
python scripts/run_benchmarks.py --config configs/grid.yaml
python scripts/run_benchmarks.py --config configs/strings.yaml
python scripts/run_benchmarks.py --config configs/chemistry_small.yaml
python scripts/run_benchmarks.py --config configs/ablations.yaml
python scripts/plot_results.py --run results/toy
python scripts/plot_results.py --run results/chemistry_small
python scripts/sweep_exact.py

The implemented comparisons include uniform executable sampling, reference-process reward tilting, trajectory balance with uniform backward probabilities, DooABLe, and the exact joint law. Ablations change route cost, terminal merging, backward normalization, and prefix-value estimation. Runs save configuration, per-seed metrics, summary CSVs, and checkpoints. Plotting writes vector PDFs, PNGs, and a LaTeX table from those measurements.

Budget and property sweeps use

python scripts/run_molecular.py --config configs/molecular_public.yaml
python scripts/sweep_preferences.py --graph results/validation/chemistry_graph.json --models results/validation/property_models

Add --neural-steps 1200 to the preference sweep to train a forward policy for each setting. Consecutive settings reuse the preceding network weights.

Local catalogs and experiment scope

Edit the input paths in configs/enamine_local.yaml , then run

python scripts/run_molecular.py --config configs/enamine_local.yaml

The catalog reader accepts CSV, TSV, and gzip files with id,smiles or native Enamine_ID,SMILES columns. Reaction templates specify a parent and one reagent. The inventory schema, budget conventions, and state features are documented in docs/data.md .

The implementation materializes the reachable graph for the supplied inventory. Each training run has a fixed graph, reward specification, and route temperature. Graph expansion can dominate cost as catalogs and budgets grow. The default local-catalog configuration uses 100 input rows and fails explicitly if its graph-size cap is exceeded.

The full LIT-PCBA docking study, licensed Enamine experiments, released molecular encoders, and third-party SynFlowNet/RxnFlow/CGFlow/SyntheMol comparisons remain to be executed. Their dependencies and manuscript mappings are listed in docs/experiments.md . The included validation measurements cover the public property predictors, small reaction inventory, and discrete systems.

Checkpoints and checks

Resume an interrupted run with the same graph and settings. --steps specifies the total update count.

dooable train --graph results/demo/graph.json --steps 2000 --resume results/demo/model --output results/demo/model
python -m pytest -q

Use --initialize PATH to reuse network weights for a new reward specification on the same graph. Exact probability checks, route replay, scaffold separation, checkpoint reload, and resumed training are covered by the tests. Completed runs and measured values appear in docs/validation.md .

Citation and licenses

The manuscript is a research draft with unfilled main-benchmark tables. Citation metadata are provided in CITATION.cff . The source code uses the MIT license. Ubuntu fonts retain the Ubuntu Font Licence. Dataset sources and third-party assets retain their upstream terms, listed in docs/data.md .