Amortized Safe Active Learning (ASAL)

Pretrained models and reproduction scripts for the paper:

Amortized Safe Active Learning for Real-Time Data Acquisition: Pretrained Neural Policies from Simulated Nonparametric Functions
Cen-You Li, Marc Toussaint, Barbara Rakitsch*, Christoph Zimmer*
AISTATS 2026

The core implementation is hosted at active-learning-framework.

Citation

@inproceedings{li2026amortized,
    title={Amortized Safe Active Learning for Real-Time Data Acquisition: Pretrained Neural Policies from Simulated Nonparametric Functions},
    author={Cen-You Li and Marc Toussaint and Barbara Rakitsch and Christoph Zimmer},
    booktitle={International Conference on Artificial Intelligence and Statistics},
    year={2026},
    url={https://openreview.net/forum?id=WWrdo1tfdw}
}

---

Installation

We use conda for environment management (conda docs).

1. Clone and Install

git clone https://github.com/cenyou/active-learning-framework
cd active-learning-framework
conda env create --file environment.yml
conda activate alef
pip install -e .

2. Verify Installation

conda activate alef
pytest tests

3. Configure Paths

Edit alef/configs/paths.py:

• Set EXPERIMENT_PATH to your experiment folder (datasets will be saved to EXPERIMENT_PATH/data)
• Set PYTHON_EXECUTOR to your alef environment's Python path

Troubleshooting

OMP Error / pytest fails loading packages

If you see OMP: Error #15: Initializing libiomp5md.dll, but found libiomp5md.dll already initialized, remove the duplicate DLL:

path\to\user\.conda\envs\alef\Lib\site-packages\torch\lib\libiomp5md.dll

PyTorch and CUDA compatibility issues

Try one of the following:

pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116

or

pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113

JAX errors

conda uninstall jax jaxlib
pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

---

Pretraining Phase

Training Unconstrained AL

python ./alef/active_learners/amortized_policies/training/main_train.py --loss_config GPMI2LossConfig --dimension D --num_experiments_initial N_init N_init --num_experiments 1 T --batch_random_shorter_sequences True --policy_knows_budget True --device cuda --seed 0 1 2 --lr 0.0001

Parameters:

• D: input dimension
• N_init: number of initial data points
• T: maximum number of queries

See Appendix Table S.F.4 for recommended settings.

Loss function options:

Loss Config	Description
GPMI2LossConfig	$\mathcal{I}_{\text{mean}}$ (default)
GPEntropy2LossConfig	$\mathcal{H}$
GPEntropy1LossConfig	$\mathcal{H}_{\text{mean}}$

Training DAD Baseline

python ./alef/active_learners/amortized_policies/training/main_train.py --loss_config DADLossConfig --dimension D --num_experiments_initial N_init N_init --num_experiments T T --batch_random_shorter_sequences False --policy_knows_budget False --device cuda --seed 0 --lr 0.0001

Training ALINE Baseline

Follow the instructions at cenyou/ALINE (asal branch).

git clone https://github.com/cenyou/ALINE
cd ALINE
git checkout asal
conda create ...
...

Training Safe AL (ASAL)

python ./alef/active_learners/amortized_policies/training/main_train.py --loss_config MinUnsafeGPEntropy2LossConfig --alpha alpha --dimension D --num_experiments_initial N_init N_init --num_experiments 1 T --batch_random_shorter_sequences True --policy_knows_budget True --device cuda --seed 0 1 2 --lr 0.0001

Note: The safety likelihood tolerance $\gamma$ from the paper is denoted as alpha in the code (in percentage). For example, alpha=5 corresponds to $\gamma=0.05$.

See Appendix Table S.F.4 for recommended (D, N_init, T) settings.

Ablation loss configs:

Loss Config	Description
MinUnsafeGPEntropy1LossConfig	$\mathcal{S}{\mathcal{H}{\text{mean}}}$
SafeGPEntropy2LossConfig	$\mathcal{S}_{\mathcal{H}, \text{division}}$
SafeGPEntropy1LossConfig	$\mathcal{S}{\mathcal{H}{\text{mean}}, \text{division}}$

---

Experiments (Deployment Phase)

Prerequisites

1. Install TabPFN/ALINE Baseline

From the active-learning-framework folder:

git checkout tabpfn
conda env create --file environment_tabpfn.yml
conda activate asal-tabpfn
pip install -e .
git checkout main

2. Install Amortized GP Baseline

git clone https://github.com/cenyou/Amor-Struct-GP/
cd Amor-Struct-GP
git checkout asal
conda activate alef
pip install -e .

Download the pretrained weights:

cd $EXPERIMENT_PATH
git clone https://github.com/boschresearch/Amor-Struct-GP-pretrained-weights

Model path: EXPERIMENT_PATH/Amor-Struct-GP-pretrained-weights/main_state_dict_paper.pth

3. Download Pretrained Policies and PFNs

Clone this repository and copy the models to EXPERIMENT_PATH:

Source	Destination
./amortized_AL	EXPERIMENT_PATH/amortized_AL
./pfns	EXPERIMENT_PATH/pfns
./ALINE	EXPERIMENT_PATH/ALINE

Note: This repository uses Git LFS to store large model files. Make sure to install Git LFS and run git lfs pull after cloning to download the pretrained weights.

4. Download Datasets

Dataset	Source	Destination
Airline	Download	EXPERIMENT_PATH/data/airlines/
LGBB	Download	EXPERIMENT_PATH/data/lgbb/
Airfoil	Download	EXPERIMENT_PATH/data/airfoil/
Engine	Download	EXPERIMENT_PATH/data/engine/

Running Experiments

Using alef environment

conda activate alef

Unconstrained AL:

Experiment	Command
AAL on Sin, Branin	python ./alef/experiments/oracle_active_learning/run_main_policy_oal.py
AAL on Airline, LGBB, Airfoil	python ./alef/experiments/pool_active_learning/run_main_policy_dataset_al.py
Baselines on Sin, Branin	python ./alef/experiments/oracle_active_learning/run_main_oal.py
Baselines on Airline, LGBB, Airfoil	python ./alef/experiments/pool_active_learning/run_main_dataset_al.py

Safe AL:

Experiment	Command
ASAL on Simionescu, Townsend	python ./alef/experiments/safe_learning/run_main_policy_safe_oal.py
ASAL on LGBB, Engine, Fluid System	python ./alef/experiments/safe_learning/run_main_policy_safe_al.py
Baselines	python ./alef/experiments/safe_learning/run_main_safe_al.py

Using asal-tabpfn environment

git checkout tabpfn
conda activate asal-tabpfn

Unconstrained AL:

Experiment	Command
ALINE on Sin, Branin	python ./alef/experiments/oracle_active_learning/run_main_aline_oal.py
ALINE on Airline, LGBB	python ./alef/experiments/pool_active_learning/run_main_aline_dataset_al.py
TabPFN on Sin, Branin	python ./alef/experiments/oracle_active_learning/run_main_oal.py
TabPFN on Airline, LGBB, Airfoil	python ./alef/experiments/pool_active_learning/run_main_dataset_al.py

Safe AL:

Experiment	Command
TabPFN Baseline	python ./alef/experiments/safe_learning/run_main_safe_al.py

---

License

See LICENSE for details.