SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1. Cubic: This setting evaluates methods under model mismatch scenario where the model being learned suffers with severe approximation error. In this task, it is important for methods to allocate model capacity to the points more critical for the downstream tasks. Prediction Model: A linear prediction model  $f_{\theta}(x) \coloneqq \theta x$  is learned for the problem where the ground truth data is generated by cubic function, i.e.,  $y_{i} = 10x_{i}^{3} - 6.5x_{i}, x_{i} \in U[-1,1]$ . Downstream task: Choose top  $B = 1$  out of  $M = 50$  resources  $\hat{\mathbf{y}} = [\hat{y}_1, \dots, \hat{y}_M]$ ,  $z^{\star}(\hat{\mathbf{y}}) \coloneqq \arg \max_i \hat{\mathbf{y}}$
Table 2: Normalized test decision quality (DQ) on the decision oriented learning problems.
Method	α	Problems
Cubic	Budget	Portfolio
MSE	-0.96±0.02	0.54±0.17	0.33±0.03
DFL	0	0.61±0.74	0.91±0.06	0.25±0.02
DFL	10	0.62±0.74	0.81±0.11	0.34±0.03
LODL	0	0.96±0.005	0.84±0.105	0.17±0.05
LODL	10	-0.95±0.005	0.58±0.14	0.30±0.03
TaskMet	0.96±0.005	0.83±0.12	0.33±0.03
0=random model 1=oracle model
Table 3: Test prediction errors (MSE) on the decision oriented learning problems.
Method	α	Problems
Cubic	Budget (×1e-4)	Portfolio (×1e-4)
MSE	2.30±0.03	4.32±2.35	4.03±0.24
DFL	0	2.89±0.32	71.7±58.3	8.0e3±8e2
DFL	10	2.41±0.05	8.09±12.1	5.18±0.46
LODL	0	2.88±0.03	35.9±12.9	55.6±9.95
LODL	10	2.29±0.19	5.05±1.88	4.31±0.31
TaskMet	2.89±0.03	9.74±13.79	4.69±0.56
$\alpha$  is the prediction loss weight in eq. (1)
2. Budget Allocation: Choose top  $B = 2$  websites to advertise based on Click-through-rates (CTRs) predictions of  $K$  users on  $M$  websites.
Prediction Model:  $\hat{\mathbf{y}}_m = f_\theta(x_m)$  where  $x_m$  is given features of  $m^{\text{th}}$  website and  $\hat{\mathbf{y}}_m = [\hat{y}_{m,1}, \dots, \hat{y}_{m,K}]$  is the predicted CTRs for  $m^{\text{th}}$  website for all  $K$  users.
Downstream task: Determine  $B = 2$  websites such that the expected number of users that click on the ad at least once is maximized, i.e.,  $z^{\star}(\hat{\mathbf{y}}_m) = \arg \max_z\sum_{j = 0}^{K}(1 - \prod_{i = 0}^M z_i\cdot \hat{y}_{ij})$  where  $z_{i}\in \{0,1\}$ .
3. Portfolio Optimization: The task is to choose a distribution over  $M$  stocks in Markowitz portfolio optimization [Markowitz and Todd, 2000, Michaud, 1989] that maximizes the expected return under the risk penalty.
Prediction Model: Given the historical data  $x_{m}$  about a stock  $m$ , predict the future stock price  $\hat{y}_{m}$ . Combining prediction over  $M$  stocks to get  $\hat{\mathbf{y}} = [\hat{y}_1, \dots, \hat{y}_M]$ .
Downstream Task: Given the correlation matrix  $Q$  of the stocks, choose a distribution over stocks  $\mathbf{z}^{\star}(\hat{\mathbf{y}}) = \arg \max_{\mathbf{z}}\mathbf{z}^{\top}\hat{\mathbf{y}} -\lambda \mathbf{z}^{\top}Q\mathbf{z}$  s.t.  $\sum_{i = 0}^{M}z_{i}\leq 1$  and  $0\leq z_{i}\leq 1,\forall i$
We run our own experiments for LODL [Shah et al., 2022] using their public code.
4.2.2 Experimental results
Table 2 presents a summary of the performance of different methods on all the tasks. Each problem poses unique challenges for the methods. The cubic setting suffers from severe approximation errors, hence the learning method needs to allocate limited model capacity more towards higher utility points compared to lower utility points. The MSE method performs the worst as it lacks task information and only care about prediction error. DFL with  $\alpha = 0$  performs better than MSE, but it can get trapped in local optima, leading to higher variance in the problem [Shah et al., 2022]. LODL  $(\alpha = 0)$  performs among the highest in this problem since it uses learned loss for each point. TaskMet performs as good as LODL as it can capture the relative importance of higher utility points versus lower utility points using the learned metric, resulting in more accurate predictions for those points (see fig. 4). In budget allocation, DFL (with  $\alpha = 0$ ) performs the best, since it is solely optimizer over  $\mathcal{L}_{\mathrm{task}}$ , but on the other hand it has 10 orders of larger prediction error as shown in table 3 indicating that the model is overfit to the task, LODL  $(\alpha = 0)$  suffers from the same problem. TaskMet has the  $2^{\mathrm{nd}}$  best Decision
Quality without overfitting on the task, i.e., low prediction error. In Portfolio Optimization, the decision quality correlates highly with the model accuracy/prediction error as in this setting the optimization problem mostly depends upon the accurate prediction of the stocks. This is the reason that MSE, DFL  $(\alpha = 10)$  performs the best, but DFL  $(\alpha = 0)$  performs the worst, since it has solely
Figure 4: (Cubic problem) TaskMet learns a metric that prioritizes points that are the most important the downstream task. The euclidean metric (MSE) puts equal weight on all points and leads to a bad model with respect to the downstream task.
Figure 5: OMD [Nikishin et al., 2022] uses the planning task loss to learn the model parameters using implicit gradients. TaskMet add one more optimization step over OMD and instead of learning the model parameters using task loss, we learn the metric which then is used to learn model parameters.
Figure 6: Results on the cartpole with distracting states [Nikishin et al., 2022].
Figure 7: Our learned metric successfully distinguishes the real states from the distracting states, i.e. the real states take a higher metric value.
being trained on  $\mathcal{L}_{\mathrm{task}}$  without any  $\mathcal{L}_{\mathrm{pred}}$ . As shown in table 2 and table 3, TaskMet is the only method that consistently performs well considering both task loss and prediction loss, across all the problem settings, this is due to the ability of the metric to infer problem-specific features without manual tuning, unlike other methods.
4.3 Model Based Reinforcement Learning
4.3.1 Background and experimental setup
Model-based RL suffers from objective-mismatch [Bansal et al., 2017, Lambert et al., 2020]. This is because dynamics models trained for data likelihood maximization do not translate to optimal policy. To reduce objective-mismatch, different losses [Farahmand et al., 2017, Voelcker et al., 2022] have been proposed to learn the model which is better suited to learning optimal policies. TaskMet provides an alternative approach towards reducing objective-mismatch, as the prediction loss is directly learnt using task loss. We set up the MBRL problem as follows. Given the current state  $s_t$  and control  $a_t$  at a timestep  $t$  of a discrete-time MDP, the dynamics model predicts the next state transition, i.e.  $\hat{s}_{t+1} := f_\theta(s_t, a_t)$ . The prediction loss is commonly the squared error loss  $\mathbb{E}_{s_t, a_t, s_{t+1}} \| s_{t+1} - f_\theta(s_t, a_t) \|^2_2$ , and the downstream task is to find the optimal Q-value/policy. Nikishin et al. [2022] introduces idea of optimal model design (OMD) to learn the dynamics model end-to-end with the policy objective via implicit differentiation. Let  $Q_\omega(s, a)$  be the action-conditional value function parameterized by  $\omega$ . The Q network is trained to minimize the Bellman error induced by the model  $f_\theta$ :
$$
\mathcal {L} _ {Q} (\omega , \theta) := \mathbb {E} _ {s, a} \left[ Q _ {w} (s, a) - \mathrm {B} ^ {\theta} Q _ {\bar {w}} (s, a) \right] ^ {2}, \tag {9}
$$
where  $\bar{\omega}$  is moving average of  $\omega$  and  $\mathrm{B}^{\theta}$  is the model-induced Bellman operator  $\mathrm{B}^{\theta}Q_{\bar{w}}(s,a)\coloneqq r_{\theta}(s,a) + \gamma \mathbb{E}_{p_{\theta}(s,a,s')}[\log \sum_{a'}\exp Q(s',a')]$ . Q-network optimality defines  $\omega$  as an implicit function of the model parameters  $\theta$  as  $\omega^{\star}(\theta) = \arg \min_{\omega}\mathcal{L}_Q(\omega ,\theta)\Rightarrow \frac{\partial\mathcal{L}_Q(\omega,\theta)}{\partial\omega} = 0$ . Now we have task loss which is optimized to find optimal Q-values:
$$
\mathcal {L} _ {\text {t a s k}} \left(\omega^ {\star} (\theta)\right) := \mathbb {E} _ {s, a} \left[ Q _ {\omega^ {\star} (\theta)} (s, a) - \mathrm {B} Q _ {\bar {\omega}} (s, a) \right] ^ {2} \tag {10}
$$
where the Bellman operator induced by ground-truth trajectory and reward is  $\mathrm{B}Q_{\bar{\omega}}(s,a)\coloneqq r(s,a) + \gamma \mathbb{E}_{s,a,s'}\log \sum_{a'}\exp Q_{\bar{\omega}}(s',a')$
Figure 8: Results on cartpole with a reduced model capacity from Nikishin et al. [2022].
OMD setup. OMD end-to-end optimizes the model for the task loss, i.e.  $\theta^{\star} = \arg \min_{\theta}\mathcal{L}_{\mathrm{task}}(\omega^{\star}(\theta))$
TaskMet setup. For metric learning, we extend OMD to learn a metric using task gradients, to train the model parameters, see fig. 5. Metric learning just adds one more level of optimization to OMD and results in the tri-level problem
$$
\phi^{\star} = \operatorname *{arg  min}_{\phi}\mathcal{L}_{\text{task}}(\omega^{\star})
$$
$$
\text {s u b j e c t} \omega^ {\star} \left(\theta^ {\star}\right) = \underset {\omega} {\arg \min } \mathcal {L} _ {Q} \left(\omega , \theta^ {\star}\right) \tag {11}
$$
$$
\theta^ {\star} (\phi) = \arg \min  _ {\mathcal {L} _ {\text {p r e d}}} (\phi , \theta)
$$
where  $\mathcal{L}_{\mathrm{task}}(\omega^{\star})$  and  $\mathcal{L}_Q(\omega ,\theta^\star)$  are defined in eq. (10) and eq. (9), respectively, and  $\mathcal{L}_{\mathrm{pred}}(\phi ,\theta) = \mathbb{E}_{s_t,a_t,s_{t + 1}}\| s_{t + 1} - f_\theta (s_t,a_t)\|_{\Lambda_\phi (s_t)}^2$  is the metricized prediction loss in eq. (3).
To learn  $\phi^{\star}$  using gradient descent, we estimate  $\nabla_{\phi} \mathcal{L}_{\text {task }}$  as
$$
\begin{array}{l} \nabla_ {\phi} \mathcal {L} _ {\text {t a s k}} = \nabla_ {\omega} \mathcal {L} _ {\text {t a s k}} \left(\omega^ {\star}\right) \cdot \frac {\partial \omega^ {\star}}{\partial \theta^ {\star}} \cdot \frac {\partial \theta^ {\star}}{\partial \phi} \\ = \nabla_ {\omega} \mathcal {L} _ {\text {t a s k}} \left(\omega^ {\star}\right) \cdot \left(\frac {\partial^ {2} \mathcal {L} \left(\omega , \theta^ {\star}\right)}{\partial \omega^ {2}}\right) ^ {- 1} \cdot \frac {\partial^ {2} \mathcal {L} \left(\omega , \theta^ {\star}\right)}{\partial \theta \partial \omega} \Bigg | _ {\omega^ {\star} \left(\theta^ {\star}\right)} \cdot \left(\frac {\partial^ {2} \mathcal {L} _ {\text {p r e d}} (\theta , \phi)}{\partial \theta^ {2}}\right) ^ {- 1} \cdot \frac {\partial^ {2} \mathcal {L} _ {\text {p r e d}} (\theta , \phi)}{\partial \phi \partial \theta} \Bigg | _ {\theta^ {\star} (\phi)} \tag {12} \\ \end{array}
$$
4.3.2 Experimental results
We replicated experiments from Nikishin et al. [2022] on the Cartpole environment. The first experiment involved state distractions, where the state of the agent was augmented with noisy and uninformative values. In this setting, we considered an unconditional diagonal metric of dimension  $n$ , which is the dimension of the state space, i.e.  $\Lambda_{\phi} \coloneqq \mathrm{diag}(\phi)$ , where  $\phi \in \mathbb{R}^n$ . As shown in fig. 6, the MLE method performed the worst across different numbers of distracting states, as it allocated its capacity to learn distracting states as well. TaskMet outperformed the other methods in all scenarios. The superior performance of TaskMet with distracting states can be attributed to the metric's ability to explicitly distinguish informative states from noise states using the task loss and then train the model using the given metric, as shown in fig. 7. The learned metric in fig. 7 assigned the highest weight to the third dimension of the state space, which corresponds to the pole angle — the most indicative dimension for the reward. This shows that the metric can differentiate state dimensions based on their importance to the task.
We also consider a setting with reduced model capacity, where the network is under-parametrized, forcing the model to prioritize how it allocates its capacity. In this scenario, we employ a full conditional metric, denoted as  $\Lambda_{\phi} = \Lambda_{\phi}(x)$ , which enables the metric to weigh dimensions and state-action pairs differently. We conducted the experiment using a model size of 3 hidden units in the layer. As depicted in fig. 8, TaskMet achieves a better return on evaluation compared to MLE and OMD. Additionally, it is evident that TaskMet achieves a lower MSE on the model predictions compared to OMD, indicating that learning with the metric contributes to a better dynamics model.
In conclusion, this paper addresses the challenge of combining task and prediction losses in task-based model learning. While task-based learning methods offer the advantage of discovering task-relevant features and data samples without explicit inductive biases, the current trend of using task loss alongside prediction loss has potential limitations. These limitations include overfitting of the prediction model to a specific task, rendering it ineffective for other tasks, and the lack of interpretability in the task-relevant features learned by the prediction model.
To overcome these limitations, the paper introduces the concept of task-driven metric learning, which integrates the task loss into a parameterized prediction loss. This approach enables end-to-end learning of metrics to train prediction models, allowing the models to focus on task-relevant features and dimensions in the prediction space. Moreover, the resulting prediction models become more interpretable, as metric learning serves as a preconditioning step for gradient-based model training. The effectiveness of the method is shown using different scales of experimental setting - decision oriented tasks as well as downstream learning tasks.
One of the limitations of the method is stability of learning the metric. Bad gradients can lead collapsed metric which can lead to unrecoverable bad predictions. Hence, hyper-parameter tuning of learning rate for metric learning and parameterization choices of the metric are crucial. Possible extensions to this work includes end-to-end metric learning with multiple task losses, learning metric for training dynamics models to be used for long-horizon planning tasks, etc.
We would like to thank Brian Karrer, Claas Voelcker, Karen Ullrich, Leon Bottou, Maximilian Nickel, and Mike Rabbat insightful comments and discussions.
Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances in neural information processing systems, 32, 2019. Cited on page 6.
Brandon Amos, Ivan Jimenez, Jacob Sacks, Byron Boots, and J Zico Kolter. Differentiable mpc for end-to-end planning and control. Advances in neural information processing systems, 31, 2018. Cited on page 2.
Anonymous. Predict-then-optimize via learning to optimize from features. In Submitted to The Twelfth International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=jev0vJ3XSjK. under review. Cited on page 2.
Somil Bansal, Roberto Calandra, Ted Xiao, Sergey Levine, and Claire J Tomlin. Goal-driven dynamics learning via bayesian optimization. In 2017 IEEE 56th Annual Conference on Decision and Control (CDC), pages 5168-5173. IEEE, 2017. Cited on page 8.
Yoshua Bengio. Using a financial training criterion rather than a prediction criterion. International journal of neural systems, 8(04):433-443, 1997. Cited on page 2.
Mohak Bhardwaj, Byron Boots, and Mustafa Mukadam. Differentiable gaussian process motion planning. In 2020 IEEE international conference on robotics and automation (ICRA), pages 10598-10604. IEEE, 2020. Cited on page 2.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. Advances in neural information processing systems, 35:5230-5242, 2022. Cited on page 5.
Tsai-Hsuan Chung, Vahid Rostami, Hamsa Bastani, and Osbert Bastani. Decision-aware learning for optimizing health supply chains. arXiv preprint arXiv:2211.08507, 2022. Cited on page 2.
Ulisse Dini. Analisi infinitesimale. Lithografia Gorani, 1878. Cited on page 14.
Bryan Wilder, Bistra Dilkina, and Milind Tambe. Melding the data-decisions pipeline: Decision-focused learning for combinatorial optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 1658-1665, 2019. Cited on pages 1, 2, 5, and 6.
Ruihan Wu, Chuan Guo, Awni Hannun, and Laurens van der Maaten. Fixes that fail: Self-defeating improvements in machine-learning systems. Advances in Neural Information Processing Systems, 34:11745-11756, 2021. Cited on page 3.
Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable top-k with optimal transport. Advances in Neural Information Processing Systems, 33:20520-20531, 2020. Cited on page 6.
Liu Yang and Rong Jin. Distance metric learning: A comprehensive survey. Michigan State University, 2(2):4, 2006. Cited on page 3.
Arman Zharmagambetov, Brandon Amos, Aaron Ferber, Taoan Huang, Bistra Dilkina, and Yuandong Tian. Landscape surrogate: Learning decision losses for mathematical optimization under partial information. arXiv preprint arXiv:2307.08964, 2023. Cited on page 3.
A The implicit function theorem
We used the implicit function theorem to compute the derivative of the prediction model with respect to the metric's parameters in eq. (7). For completeness, this section briefly presents the standard implicit function theorem, cf. Dini [1878] and Dontchev and Rockafellar [2009, Theorem 1B.1]:
Theorem 1 (Implicit Function Theorem) Implicit Function Theorem: Let  $f:\mathbb{R}^n\times \mathbb{R}^m\to \mathbb{R}$  be a continuous differentiable function, and let  $x^{\star},y^{\star}$  be a point satisfying  $f(x^{\star},y^{\star}) = 0$ . If the Jacobian  $\frac{\partial(f(x^{\star},y^{\star})}{\partial y}\neq 0$ , then there exists an open set around  $(x^{\star},y^{\star})$  and a unique continuously differentiable function  $g$  such that  $y^{\star} = g(x^{\star})$  and  $f(x,g(x)) = 0$ . Additionally, the following relation holds:
$$
\frac {\partial g (x)}{\partial x} = - \left(\frac {\partial f (x , y ^ {\star})}{\partial y}\right) \frac {\partial f (x , y ^ {\star})}{\partial x} | _ {y ^ {\star} = g (x)} \tag {13}
$$
B.1 Decision Oriented Model Learning
We replicated our experiments using the codebase provided by Shah et al. [2022], which can be found on github. To ensure consistency, we used the same hyperparameters as mentioned in the code or article for the baselines. Our metric learning pipeline was added on top of their code, and thus we focused on tuning hyperparameters related to metric learning. The metric is parameterized as  $\Lambda_{\phi}(x) = L_{\phi}(x)L_{\phi}^{\top}(x) + \epsilon_{\phi}\mathbb{I}_{n\times n}$ , where  $\epsilon_{\phi}$  is a learnable parameter that explicitly controls the amount of Euclidean metric in the predicted metric. This helps ensure the stability of metric learning. We initialize the parameters in such a way that the predicted metric is close to the Euclidean metric. For each outer loop of metric update, we perform  $K$  inner updates to train the predictor. Following the methodology of Shah et al. [2022], we conducted 50 runs with different seeds for each of the experiments, where each method was evaluated on 10 different datasets, with 5 different seeds used for each dataset.
Table 4: Hyper-parameters for Decision Oriented Learning Experiments
Hyper-Parameter	Values
\( \Lambda_{\phi} \) learning rate	\( 10^{-3} \)
\( \Lambda_{\phi} \) hidden layer sizes	[200]
Warmup steps	500
Inner Iterations (K)	100
Implicit derivative batchsize	10
Implicit derivative solver	Conjugate gradient on the normal equations (5 iterations)
B.2 Model-Based Reinforcement Learning
We consider the work of Nikishin et al. [2022] as the baseline for replicating the experiments, and we build upon their source code. Our metric learning is just one additional step to their method. We adopt exact same hyperparameters as their for dynamics learning and Action-Value function learning. We focus on exploring and tuning the hyper-parameters specific to the metric learning component of the method.
Table 5: Hyper-parameters for the CartPole experiments
Hyper-Parameter	Values
\( \Lambda_{\phi} \) learning rate	\( 10^{-3} \)
\( \Lambda_{\phi} \) hidden layer sizes	[32, 32]
Warmup steps	5000
Inner iterations (K)	1
Implicit derivative batchsize	256
Implicit derivative solver	Conjugate gradient on the normal equations (10 iterations)
For the state distractor experiments, we parameterize the metric as an unconditional diagonal matrix, denoted as  $\Lambda_{\phi} = \mathrm{diag}(\phi)$  where  $\phi \in \mathbb{R}^n$  and  $n$  is the dimension of the state space. In addition, we consider a hyper-parameter of metric parameterization, for which we either take normalize or unnormalized metric. When we refer to normalizing the metric, we mean constraining the norm of the  $\phi$  vector to be equal to the L2 norm of an euclidean metric which is used by MSE method. This constrains the family of learnable metrics. To achieve this, we set  $\phi := \frac{\phi}{\|\phi\|_2}\sqrt{n}$ , ensuring  $\| \phi \|_2 = \| \mathbb{I}_{n\times n}\| _2 = \sqrt{n}$ . We also used L1 regularization on the metric output, to induce sparsity in the metric. We sweep over three values of the regularization coefficient - [0.0, 0.001, 0.1]. We ran a sweep over the 6 combinations of hyperparameters - [unnormized, normalized]  $\times$  [0.0, 0.001, 0.1] and choose the best hyper-parameter combination for each of the experiment. All the number reported in the experiments are calculated over 10 random seeds.
Our metric learning approach uses two implicit gradient steps. Firstly, we take the implicit derivative through action-value network parameters, approximating the inverse hessian to the identity, similar to Nikishin et al. [2022]. Secondly, for the step through dynamics network parameters, we calculate the exact implicit derivative.