SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dose-response estimator: For training our DCNet  $\hat{\mu}(t,x)$ , we follow previous works from the baseline VCNet [2, 56]. We set the representation network  $\phi$  to a multi-layer perceptron (MLP) with two hidden layers with 50 hidden neurons each and ReLU activation function. For the parameters of the prediction head  $h$ , we use the same model choices as for  $\phi$ . Additionally, we use B-splines with degree two and place the knots of the tensor product basis at  $\{1/3,2/3\}^p$ . For the baseline MLP from Sec. 5, we ensure similar flexibility for fair comparison, and we thus select a MLP with four hidden layers with 50 hidden units each and ReLU activation. We train the networks minimizing the mean squared error (MSE) loss  $\mathcal{L}_{\mu} = \frac{1}{n}\sum_{i=1}^{n}\left(\hat{\mu}(t_i,x_i) - y_i\right)^2$ . For optimization, we use the Adam optimizer [43] with batch size 1000 and train the network for a maximum of 800 epochs using early stopping with a patience of 50 on the MSE loss on the factual validation dataset.
We tune the learning rate within the search space  $\{0.0001, 0.0005, 0.001, 0.005, 0.01\}$ , and, for evaluation, we use the same criterion as for early stopping.
Conditional normalizing flows: For modeling the GPS  $\hat{f}(t, x)$ , we use conditional normalizing flows (CNFs). Specifically, we use neural spline flows [15] in combination with masked auto-regressive networks [13]. We use a flow length of 1 with 5 equally spaced bins and quadratic splines. For the autoregressive network, we use a MLP with 2 hidden layers and 50 neurons each. Additionally, we use noise regularization with noise sampled from  $N(0, 0.1)$ . We train the CNFs by minimizing the negative log-likelihood (NLL) loss  $\mathcal{L}_f = -\frac{1}{n} \sum_{i=1}^{n} \log \hat{f}(t_i, x_i)$ . For optimization, we use the Adam optimizer with batch size 512 and train the CNFs for a maximum of 800 epochs using early stopping with a patience of 50 on the NLL loss on the factual validation dataset.
We tune the learning rate within the search space  $\{0.0001, 0.0005, 0.001, 0.005, 0.01\}$  and for evaluation we use the same criterion as for early stopping.
Policy learning: After training the nuisance models  $\hat{\mu}$  and  $\hat{f}$ , we use them as plug-in estimators for our policy learning. For the policy network  $\hat{\pi}_{\theta}$ , we select a MLP with 2 hidden layers with 50 neurons each and ReLU activation. We set the reliability threshold  $\bar{\varepsilon}$  to the  $5\%$  -quantile of the estimated GPS  $\hat{f}(t,x)$  of the train set, when not specified otherwise. We use Adam optimizers for updating  $\theta$  and  $\lambda$  each, with batch size 512, and train the network using the gradient descent-ascent optimization objective from Eq. (7) for a maximum of 400 epochs using early stopping with a patience of 20 on the validation loss from Eq. (9) on the factual validation dataset.
We set the learning rate for updating  $\lambda$  to  $\eta_{\lambda} = 0.01$  and use random search with 10 configurations to tune the learning rate for updating the parameters of the policy network  $\eta_{\theta} \in \{0.0001, 0.0005, 0.001, 0.005, 0.01\}$  and the initial value of the Lagrangian multipliers  $\lambda_i \in [1, 5]$ . To evaluate the performance during hyperparameter search, we use the same criterion as for early stopping. Then, after determining the hyperparameters, we perform  $k = 5$  runs to find the optimal policy.
Our above procedure was chosen to ensure a fair and direct comparison to previous literature and to the naive baselines, which enables us to show the sources of performance gain in our method (see Sec. 5). Note that when applying our method for reliable policy learning for dosage combinations in production, one might consider optimizing more exhaustively also over the different modeling choices of the nuisance estimators and the policy network.