Algorithm 1 Training algorithm of APID
1: Input. Counterfactual query  $Q_{a^{\prime}\to a}(y^{\prime})$ ; training dataset  $\mathcal{D} = \mathcal{D}_0\cup \mathcal{D}_1$ ; number of training iterations  $n_B,n_Q,n_{CQ}$ ; minibatch size  $b$ ; learning rates  $\eta$ ; intensity of the noise regularization  $\sigma^2$ ; the variance of the augmentation  $\varepsilon^2$ ; EMA smoothing  $\gamma$ ; loss coefficients  $\lambda_{Q}$  and  $\lambda_{\kappa}$ .  
2: Init. Flows parameters: residual NFs  $(\beta_0,\beta_1)$  and variational augmentations  $(\theta_0,\theta_1)$ .  
4: for  $k = 0$  to  $n_B$  do ▷Burn-in stage  
5: for  $a\in \{0,1\}$  do  
6:  $\{A_i = a,Y_i\}_{i = 1}^b\gets$  minibatch of size  $b$  from  $\mathcal{D}_a$   
7:  $\xi_i\sim N(0,\sigma^2);\tilde{Y}_i\gets Y_i + \xi_i$  ▷ Noise regularization  
8:  $\hat{Y}_{\mathrm{aug},i}\sim N(g^{a}(\hat{Y}_{i}),\varepsilon^{2})$  ▷ Variational augmentations  
9:  $\mathcal{L}_{\mathrm{NLL}}(\beta_a,\theta_a)\gets -\frac{1}{n}\sum_{i = 1}^{b}\log \hat{\mathbb{P}}_{\beta_a,\theta_a}(Y = \tilde{Y}_i|a)$  ▷ (1) Negative log-likelihood loss  
10:  $\{(U_{Y^{1},j},U_{Y^{2},j})\}_{j = 1}^{b}\sim \mathrm{Unif}(0,1)^{2}$   
11:  $\{(\hat{Y}_{\mathrm{aug},a,j},\hat{Y}_{a,j})\}_{j = 1}^{b}\gets \hat{F}_a\left(\{(U_{Y^{1},j},U_{Y^{2},j})\}_{j = 1}^{b}\right)$   
12:  $\mathcal{L}_{\mathbb{W}}(\beta_a)\gets \int_0^1 |\hat{\mathbb{F}}_{\hat{Y}_a}^{-1}(q) - \hat{\mathbb{F}}_{Y_a}^{-1}(q)|\mathrm{d}q$  ▷ (2) Wasserstein loss  
13:  $\mathcal{L}\gets \mathcal{L}_{\mathrm{NLL}}(\beta_a,\theta_a) + \mathcal{L}_{\mathbb{W}}(\beta_a)$   
14: optimization step for  $\beta_{a},\theta_{a}$  wrt.  $\mathcal{L}$  with learning rate  $\eta$   
end for  
end for  
End for  
Output. Pretrained flows parameters: residual NFs  $(\beta_0,\beta_1)$  and variational augmentations  $(\theta_0,\theta_1)$   
18:  
19:  $\overline{\beta}_a,\underline{\beta}_a\gets \beta_a;\quad \overline{\theta}_a,\underline{\theta}_a\gets \theta_a;$  ▷ Copying counterfactual flow parameters  
20:  $[\hat{l}_a,\hat{u}_a]\gets [\min_i(Y_i);\max_i(Y_i)]$   
21: for  $k = 0$  to  $n_Q + n_{CQ}$  do ▷ Query & curvature-query stages  
22: for  $(\beta_a,\theta_a)\in \{(\overline{\beta}_a,\overline{\theta}_a),(\underline{\beta}_a,\underline{\theta}_a)\}$  do  
23:  $\{A_i = a,Y_i\}_{i = 1}^b\gets$  minibatch of size  $b$  from  $\mathcal{D}_a$   
same procedure as in lines 7-12 to evaluate  $\mathcal{L}_{\mathrm{NLL}}(\beta_a,\theta_a)$  and  $\mathcal{L}_{\mathbb{W}}(\beta_a)$   
25:  $\mathcal{L}\gets \mathcal{L}_{\mathrm{NLL}}(\beta_a,\theta_a) + \mathcal{L}_{\mathbb{W}}(\beta_a)$   
26:  $(y',\{Y_{\mathrm{aug},j}'\}_{j = 1}^{b})\sim N(g^{a'}(y'),\varepsilon^{2}))$  ▷ Variational augmentations  
27:  $\{(\tilde{u}_{Y^1,j},\tilde{u}_{Y^2,j})\}_{j = 1}^b\gets \hat{F}_{a'}^{-1}\left(\{(y',Y_{\mathrm{aug},j}')\}_{j = 1}^b\right)$   
28:  $\{(\hat{Y}_{\mathrm{aug},j},\hat{Y}_j)\}_{j = 1}^b\gets \hat{F}_a\left(\{(\tilde{u}_{Y^1,j},\tilde{u}_{Y^2,j})\}_{j = 1}^b\right)$   
29:  $\hat{Q}_{a'\to a}(y') = \frac{1}{b}\sum_{j = 1}^{b}\hat{Y}_j$  ▷ Estimated ECOU [ECOT]  
30: if  $\hat{Q}_{a'\to a}(y')$  not in  $[\hat{l}_a,\hat{u}_a]$  and  $k < n_Q$  then continue ▷ Non-informative bounds  
32: else  
33:  $\mathcal{L}_Q(\beta_{a'},\theta_{a'},\beta_a)\gets$  Softplus(=hat $Q_{a'\to a}(y')$  ▷ (3) Counterfactual query loss  
34:  $\mathcal{L}\gets \mathcal{L} + \lambda_Q\mathcal{L}_Q(\beta_{a'},\theta_{a'},\beta_a)$   
end if  
end for  
end for  
end for  
end for