Reproduction code for "Addressing Instrument-Outcome Confounding in Mendelian Randomization through Representation Learning"

Code accompanying "Representation Learning for Mendelian Randomization" (ICML 2026 camera-ready).

This folder reproduces all simulation figures and tables in the main paper and appendix.

Environment

pyenv virtualenv 3.10 rl4mr
pyenv activate rl4mr
pip install -r requirements.txt            # minimal pinned set
# or, for an exact reproduction:
pip install -r full_freeze.txt

Key dependencies: PyTorch 2.7.1, Lightning 2.5, Hydra 1.3, ivmodels 0.6, wandb, doubleml.

Running

Single training run

python train.py <hydra overrides>

Hydra config root: conf/config.yaml. Encoder/decoder configs in conf/encoder/, conf/decoder/.

SLURM sweeps

run_grid.sh is the array-job driver. The submit_*.sh scripts set environment variables (mixing function, encoder type, loss weights, etc.) and submit run_grid.sh across data seeds ds42–ds61. See the submit-script table below for which script drives which figure/table.

Aggregation

bash run_summarize_seeds.sh    # → summarize_seeds.py + make_boxplot.py
python evaluate.py             # post-training IV estimation (TSLS) + MR-Egger baseline

What produces what

Figures and tables below are referenced by their LaTeX labels in the camera-ready manuscript (the printed numbering depends on how the document is compiled, so labels are the stable identifier).

Simulation figures

All figures are produced by make_plots.py; PNGs are written to results/writeup_simu_fig*.png.

File	Paper label	Description
writeup_simu_fig1.png	Figure 4	Bias comparison under different mixing functions (polynomials of degree 1–3 and invertible MLP)
writeup_simu_fig2.png	Figure 5	Estimation bias under misspecified latent dimensions (true p=2, varying p̂)
writeup_simu_fig3.png	Figure 6	Estimation bias with and without the independence loss
writeup_simu_fig4.png	Figure 9 (left)	Three-population experiment — bias under different mixing functions (Appendix E.3)
writeup_simu_fig5.png	Figure 9 (right)	Three-population experiment — bias under different invariance losses, polynomial mixing of degree 3 (Appendix E.3)
writeup_simu_fig6.png	Figure 10	Estimation bias as a function of an added relatedness loss weight λ₃ ∈ {0, 1e-3, 1e-2, 1e-1} (Appendix E.3)

Appendix tables

Produced by make_tables.py; CSVs are written to results/rebuttal_table*_*.csv.

Paper label	Description	Producing function
Table 1	3×3 mean-bias matrix for 2SLS(Ŵ): independence loss × invariance loss, polynomial mixing of degree 3	table1_kernels
Table 2	3×3 mean-bias matrix for PO(V̂)-2SLS(Ŵ): same axes as above	table1_kernels
Table 3	Final invariance loss (MMD on Ŵ) for the Figure 5 experiment, mean/min/max across seeds for each p̂ ∈ {1,2,3,4}	table3b_misspec_mmd

make_tables.py also defines table2_iv_methods, table3a_misspec_gcm, and table4_lam3 — these produced supplementary outputs (IV-estimator comparison, GCM rejection rates, λ₃-as-table) that did not make it into the camera-ready, and can be ignored or removed.

Submit scripts → which experiments they drive

Script	Role
submit_fixmix_inv.sh	Main workhorse. Sweeps mixing functions (cayleypoly / normalsmoothpoly / normalclamppoly / invmlp), invariance losses (meanvar / poly2 / poly3 / rbf), independence losses, and λ₃ values. Drives Figure 4, Figure 6, Figure 9 (right), Figure 10, and all three appendix tables. Set MIX_TYPE, INV_LOSSES, etc. at the top of the script
submit_fixmix_invmlp.sh	Invertible-MLP mixing branch in Figure 4, Figure 6, Figure 9 (right)
submit_fixmix_hvhw.sh	Misspecified latent dim p̂ ∈ {1,2,3,4} — Figure 5 and Table 3
submit_fixmix_3pop.sh	Three-population experiments — Figure 9

Outputs land in outputs/${EXP_NAME}-ds${DATA_SEED}/${SIM_ID}/. After training, run bash run_summarize_seeds.sh (or call summarize_seeds.py directly) to produce per-experiment summary CSVs under results/${EXP_GRP}/summary_${EXP_NAME}_*.csv, which make_plots.py and make_tables.py then consume.

File overview

File	Purpose
train.py	Lightning training entrypoint (Hydra)
mdcrl.py	LitAutoEncoder model, AdaptiveLossBalancer, dataset
torch_losses.py	MMD / HSIC kernels (mean-var, RBF, Poly 2, Poly 3, orthogonality)
simu_data.py, data_setup.py	Multi-population interventional DGP
poly_decoder.py, invert_poly_encoder.py	Polynomial decoder / oracle encoder
evaluate.py, utils_eval.py	Post-training IV estimation (TSLS), GCM test, MR-Egger baseline
competitors.py	MREgger baseline
summarize_seeds.py, separate_runs.py	Aggregate metrics across seeds; split outputs
make_plots.py	All main paper figures
make_tables.py	All Appendix tables and other tables in rebuttal
make_boxplot.py	Quick boxplots from summarized CSVs
entropy_estimators.py, hsic.py, poly_mixing.py, utils_train.py	Supporting utilities