TaskMet: Task-Driven Metric Learning for Model Learning

Dishank Bansal 1 Ricky Chen 1 Mustafa Mukadam 1 Brandon Amos 1!" # $!⋆ %!"#$%&'( "#$% )*+ ,!

Abstract!! /*Deep learning models are often used with somedownstream task. Models solely trained toachieve accurate predictions may struggle to perform well on the desired downstream tasks. We
propose using the task’s loss to learn a metric
which parameterizes a loss to train the model.
This approach does not alter the optimal prediction model itself, but rather changes the model
learning to emphasize the information important
for the downstream task. This enables us toachieve the best of both worlds: a prediction
model trained in the original prediction space
while also being valuable for the desired downstream task. We validate our approach throughexperiments conducted in two main settings: 1)
decision-focused model learning scenarios involving portfolio optimization and budget allocation,and 2) reinforcement learning in noisy environments with distracting states.

/, IFT##$%&.!⋆ ( &'()
Backprop

Figure 1. TaskMet learns a metric for parameterized prediction
loss using the gradient signal from a downstream task loss.

et al., 2022), and budget allocation, matching, and recommendation (Kang et al., 2019; Wilder et al., 2019; Shahet al., 2022). We notate the model’s loss on the prediction
problem as Lpred and the downstream task loss as Ltask .
Running examples of predictions and tasks. We focus
on the following examples in this paper: 1) the portfolio
setting from Wilder et al. (2019), which predicts the expected returns from stocks for a financial portfolio. Here,the Lpred is the MSE and Ltask is from the regret of runninga portfolio optimization problem on the output; 2) the allocation setting from Wilder et al. (2019), which predicts thevalue of items that are being allocated, e.g. click-throughrates for recommender systems. Here, Lpred is the MSE
and Ltask measures the result of allocating the highest-value
items. 3) the model-based reinforcement learning setting oflearning the system dynamics from Nikishin et al. (2022).
Here, Lpred is the MSE and the Ltask measures how well the
dynamics performs for downstream value predictions.

1. Introduction
Machine learning prediction models are typically trained tomaximize the likelihood on a training dataset. While themodels are capable of approximating the underlying data
generating process to predict the output, they are prone toapproximation errors due to limited training data and model
capacity. These errors lead to suboptimal performance indownstream tasks where the models are used. End-to-end
task-based model learning methods is an the area of machinelearning research that uses information from the downstreamtask to improve the model’s performance on that particulartask. These methods work well for financial price predictions (Bengio, 1997; Elmachtoub & Grigas, 2022), inventory
stock, demand, and price forecasting (Donti et al., 2017; Elmachtoub & Grigas, 2022; El Balghiti et al., 2019; Mandi
et al., 2020; Liu et al., 2023), dynamics modeling for RL(Farahmand et al., 2017; Amos et al., 2018; Farahmand,
2018; Bhardwaj et al., 2020; Voelcker et al., 2022; Nikishin
1
Meta AI. Correspondence to:hankbansal09@gmail.com>.

⋆"⋆

Limitations of task-based learning. Task-based model
learning comes with the promise of being able to discovertask-relevant features and data-samples on its own without
the need of explicit inductive biases. The current trend forend-to-end model learning uses task loss along with theprediction loss to train the prediction models. Though easy
to use, these methods may be limited by 1) the prediction
overfitting to the particular task, rendering it unable to generalize; 2) the need of tuning weights to combine the task
and prediction losses as in Eq. (1).
Our contributions. We present TaskMet — a task-driven
end-to-end metric learning framework which is used fortraining prediction model, see Fig. 1. This enables more interpretable learning of the model using the metric comparedto learning with combination of task loss and prediction
loss. The learned metric can uncover underlying propertiesof the task that are useful for training the model, e.g. as
in Figs. 2 and 4. Section 4 shows the empirical success of
metric learning on decision focused model learning setting
and MBRL with policy learning as the downstream task.

Dishank Bansal <dis-

Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1TaskMet: Task-Driven Metric Learning for Model Learning

2. Preliminaries and Background

Algorithm 1 TaskMet
Models: predictor fθ and metric Λϕwhile unconverged dofor i in 1 . . . K do
θ ← update(θ, ∇θ Lpred (θ, ϕ)) // Eq. (3): fit f to the current metric loss
end forϕ ← update(ϕ, ∇ϕ Ltask ) // Eq. (6): update Λ with the task loss
end while
return optimal parameters θ, ϕ

We consider a prediction model ŷ := fθ (x) : X → Y
that maps from X to Y and is parameterized by θ. The
dataset D := {(xi , yi )}N
i=1 consists of N input-output(x, y) ∈ X × Y pairs. The prediction model has an associated prediction loss, Lpred and is used in conjunction
with some downstream task that provides a task loss Ltask
characterizing how well the model performs on the task.
The most relevant related work to ours includes the approaches of Bengio (1997); Donti et al. (2017); Farahmand
et al. (2017); Kang et al. (2019); Wilder et al. (2019); Nikishin et al. (2022); Shah et al. (2022); Voelcker et al. (2022);
Nikishin et al. (2022), which learn the optimal prediction
model parameter θ to minimize the task loss Ltask :θ⋆ := arg min Ltask (θ) + αLpred (θ),

Learning model parameters with a given metric can condition the learning process in the following ways: 1) Relative
importance of dimensions: the metric allows for down- or
up-weighting different dimensions of the prediction spaceby changing the diagonal entries of the metric. 2) Correlation in the prediction space: the quadratic nature of theloss with the metric allows the model to be aware of correlations between dimensions in the prediction space. 3)
Relative importance of samples: heteroscedastic metrics
Λ(x) enable different samples to be weighted differently forthe final expected cost over the dataset. Hence, learningmodel parameters with a metricized loss can be seen as
conditioning the learning problem. The ability to learn themetric end-to-end enables the task to condition the learningof the model in any or all of the three ways described above.
This approach offers an interpretable method for the task toguide the model learning, in contrast to relying solely ontask gradients for learning model parameters, which may or
may not align effectively with the given prediction task.

(1)θ

where α is a regularization parameter to weigh the prediction
loss which is MSE error (Eq. (2)) in general. Alternatives
to Eq. (1) include 1) Smart, “Predict, then Optimize” (SPO)
methods (Elmachtoub & Grigas, 2022; El Balghiti et al.,2019; Mandi et al., 2020; Liu et al., 2023), which consider
surrogates for when the derivative is undefined or uninformative, or 2) changing the prediction space from the original
domain X into a latent domain with task information, e.g.
task-specific latent dynamics for RL (Hafner et al., 2019b;a;Hansen et al., 2022). Sadana et al. (2023) provide a broader
survey of the field of decision making under uncertainty.

3.2. End-to-end metric learning for model learning

3. TaskMet: Task-Driven Metric Learning

The key idea of the method is to learn a metric end-to-end
with a given task, which is then used to train the prediction
model as shown in Eq. (3). The learning problem of themetric and model parameters are

3.1. Metricized loss for training prediction models
A supervised prediction model fθ parameterized by θ isoften trained using the mean squared error (MSE)θ⋆ := arg min Ex,y∼D [(fθ (x) − y)2 ].

(2)

(4)

s.t. θ⋆ (ϕ) = arg min Lpred (θ, ϕ)

(5)

ϕ

θ

Equation (2) is equivalent to maximizing the data likelihoodand is often optimized using a stochastic first-order method,e.g. as in Bottou et al. (2018). MSE error is indicative of
the prediction performance of the model. We propose touse a parameterized variant of the MSE to turn it into a
metricized loss:
Lpred (θ, ϕ) := Ex,y∼D [∥fθ (x) − y∥2Λϕ (x) ],ϕ⋆ := arg min Ltask (θ⋆ (ϕ)),θ

where ϕ and θ are (respectively) the metric and model parameters, Lpred is the metricized prediction loss (Eq. (3)) totrain the prediction model, and Ltask is the task loss definedby the task at hand (which could be another optimization
problem, e.g. Eq. (8), or another learning task, e.g. Eq. (10).

(3)

Gradient-based learning. We learn the optimal metric
Λϕ⋆ with the gradient of the task loss, i.e. ∇ϕ Ltask (θ⋆ (ϕ)).
Using the chain rule and assuming we have the optimalθ⋆ (ϕ) for some metric parameterization ϕ, this derivative iswhere ∥x∥M := (x⊤ M x)1/2 is a Mahalanobis norm withpositive semi-definite matrix M , e.g. as in Ghojogh et al.
(2022). We consider a parametric metric Λϕ with parameters ϕ and make it conditional on the feature x so it canlearn the importance of the regression space from each partof the feature space. Learning θ using Lpred (θ, ϕ) makes
the model parameters θ an implicit function of ϕ denotedas θ⋆ (ϕ), with the condition ∇θ Lpred (θ, ϕ)|θ=θ⋆ = 0.

∇ϕ Ltask (θ⋆ (ϕ)) = ∇θ Ltask (θ) θ=θ⋆ (ϕ) ·∂θ⋆ (ϕ)∂ϕ(6)

To calculate the term ∇ϕ Ltask (θ⋆ (ϕ)), we need to compute
two gradient terms: ∇θ Ltask (θ) θ=θ⋆ (ϕ) and ∂θ⋆ (ϕ)/∂ϕ.
2

TaskMet: Task-Driven Metric Learning for Model Learning

The former can be estimated in standard way since Ltask (θ)
is an explicit function of θ. However, the latter cannot bedirectly calculated because θ⋆ is a function of optimization
problem which is multiple iterations of gradient descent, asshown in Eq. (5). Backpropping through multiple iterationsof gradient descent can be computationally expensive, so we
use the implicit function theorem on the first-order optimal∂L (θ,ϕ)
ity condition of Eq. (5), i.e. pred∂θ= 0; see Appendix A⋆
for more details. ∇ϕ Ltask (θ (ϕ)) can be computed with−∇θ Ltask (θ) ·∂ 2 Lpred (θ, φ)
∂θ2

!−1

∂ 2 Lpred (θ, φ)·∂φ∂θTable 1. Comparison of the normalized test decision quality(0=random, 1=oracle) on the decision oriented learning problems.
α is prediction loss weight in Eq. (1)Problems
Methodα Cubic

BudgetPortfolio

MSE
DFL
DFL
LODL
LODL

−0.96±0.02
0 0.61±0.74
10 0.62±0.74
0 0.96±0.005
10 −0.95±0.005

0.54±0.17
0.91±0.06
0.81±0.11
0.84±0.105
0.58±0.14

0.33±0.03
0.25±0.02
0.34±0.03
0.17±0.05
0.30±0.03

0.83±0.12

0.33±0.03

0.96±0.005

TaskMet(7)θ=θ ⋆ (ϕ)

3

We follow Blondel et al. (2022) and compute the implicit
derivative by using conjugate on the normal equations.

y(x) 0MSE−32
Λ(x) 10Discussion. Other related work on metric learning such as
Hastie & Tibshirani (1995); Yang & Jin (2006); Weinberger& Tesauro (2007); Kulis et al. (2013); Hauberg et al. (2012);
Kaya & Bilge (2019) often learns a non-Euclidean metricor distance that captures the geometry of the data and thensolves a prediction task such as regression, clustering, orclassification in that geometry. In contrast to these, in thetask-based model learning, we propose that the downstreamtask (instead of the data alone) gives the relevant metricfor the prediction, and that it is possible to use end-to-end
learning as in Eq. (4) to obtain the task-based metric.

−1

0

1

x

Figure 2. (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.

et al., 2019). Appendix B.2 provides the exact formulationsof g(z, y) and description for the optimization problems.
We use Eq. (4) and Eq. (5) to find optimal parameters.

4. Experiments

These settings evaluates the ability of TaskMet to capture
the correlation between model predictions and differentiate
between different data-points in accordance to their importance for the optimization problem. Hence, we consider aheteroscedastic metric i.e., Λϕ (x).

We evaluate our method on standard decision-focused model
learning and model-based reinforcement learning settings.
4.1. Decision Oriented Model Learning

Baselines. We compare our method with standard baselines
such as MSE and DFL (Donti et al., 2017; Shah et al., 2022)
along with more sophisticated method such as LODL (Shahet al., 2022). The description of baseline methods can befound in the Appendix B.2. We run our own experimentsfor LODL (Shah et al., 2022) using their public code.

Setup. We use three standard resource allocation tasks forcomparing task-based learning methods (Shah et al., 2022;
Wilder et al., 2019; Donti et al., 2017; Futoma et al., 2020).
In this setting, resource prediction based on some input
features constitutes prediction model, resource allocation
constitutes downstream task which is characterized by Ltask
The prediction model’s output is used in a downstream
resource optimization. The settings are implemented exactly
as in Shah et al. (2022) and have task losses defined by
Ltask (θ) := Ex,y∼D [g(z ⋆ (ŷ), y)]

TaskMetResults. Table 1 presents a summary of the performanceof 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 prioritize the accuracy of higher utility
points over the accuracy of lower utility points. The MSE
method performs the worst as it lacks task information andonly care about prediction error. DFL with α = 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 (α = 0) performs among the highest in this problem since it uses metricized loss as well learned for eachpoint. TaskMet performs as good as LODL as it can capture the relative importance of higher utility points versus(8)

where z ⋆ (ŷ) := arg minz g(z, ŷ) and g(z, y ′ ) is some combinatorial optimization objective over variable z parameterized by y ′ . The task loss Ltask is the expected value ofobjective function with decision variable z ⋆ (ŷ) induced by
the prediction model ŷ = fθ (x) under the ground truth
parameters y. We use corresponding surrogate losses toreplicate the z ⋆ (ŷ) optimization problem as in Shah et al.
(2022); Wilder et al. (2019); Xie et al. (2020) and differentiate through the surrogate using cvxpylayers (Agrawal
3

TaskMet: Task-Driven Metric Learning for Model Learning
real statesstate dimension
1 2 3 4 5 6 7 8 9 10 11 12 ...

500

lower utility points using the learned metric, resulting in
more accurate predictions for those points (see Fig. 2). In
budget allocation, DFL (with α = 0) performs the best,
since it is solely optimizer over Ltask , but on the other hand
it has 10 orders of larger prediction error as shown in Table 2 indicating that the model is overfit to the task, LODL(α = 0) suffers from the same problem. TaskMet has the
2nd best Decision Quality without overfitting on the task i.e.,
low prediction error. In Portfolio Optimization, the decisionquality 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 thereason that MSE, DFL (α = 10) performs the best, but DFL(α = 0) performs the worst, since it has solely being trainedon Ltask without any Lpred . As shown in Table 1 and Table 2,
TaskMet is the only method that consistently performs wellacross all the problem setting, this is due to the ability of
metric to infer problem specific features without manual
tuning unlike other methods.

300

200

100

16

32

64

128

256

metric values (darker=higher)

16
32
64
128
256
512

512

# Distractors

Figure 4. Our learned metric successfully distinguishes the realstates from the distracting states,
Figure 3. Results on the carti.e. the real states take a higher
pole with distracting states
metric value.
(Nikishin et al., 2022).
TaskMetMLE

OMD

optimization to OMD and results in the tri-level problem
φ⋆ = arg min Ltask (ω ⋆ )ϕsubject to ω ⋆ (θ⋆ ) = arg min LQ (ω, θ⋆ )
ω(11)

θ⋆ (φ) = arg min Lpred (φ, θ)θwhere Ltask (ω ⋆ ) and LQ (ω, θ⋆ ) are defined in Eq. (10) andEq. (9), respectively, and Lpred (ϕ, θ) = Est ,at ,st+1 ∥st+1 −fθ (st , at )∥2Λϕ (st ) is the metricized prediction loss in Eq. (3).
Appendix B.3 describes how we compute gradients.

4.2. Model Based Reinforcement Learning
Setup. Given the current state st and control at at atimestep t of a discrete-time MDP, the dynamics model
predicts the next state transition, i.e. ŝt+1 := fθ (st , at ).
The prediction loss is commonly the squared error loss
Est ,at ,st+1 ∥st+1 − fθ (st , at )∥22 , 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ω (s, a) be the action-conditional
value function parameterized by ω. The Q network is trained
to minimize the Bellman error induced by the model fθ :LQ (ω, θ) := Es,a [Qw (s, a) − Bθ Qw̄ (s, a)]2 ,# distractors

Episode Return400

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 wasaugmented with noisy and uninformative values to increase
dimensionality. In this setting, we considered an unconditional diagonal metric of dimension n, which is the dimension of the state space, i.e. Λϕ := diag(ϕ), where ϕ ∈ Rn .
As shown in Fig. 3, the MLE method performed the worst
across different numbers of distracting states, as it allocatedits capacity to learn distracting states as well. TaskMet
outperformed the other methods in all scenarios. The superior performance of TaskMet with distracting states canbe attributed to the metric’s ability to explicitly distinguishinformative states from noise states using the task loss andthen train the model using the given metric, as shown in
Fig. 4. The learned metric in Fig. 4 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. Appendix B.3 also shows experiments with a limited capacity.

(9)

where ω̄ is moving average of ω and Bθ is the modelθ
induced Bellman
P operator B Qw̄ (s, a) := rθ (s, a) +γ Epθ (s,a,s′ ) [log a′ exp Q(s′ , a′ )] Q-network optimality
defines ω as implicit function of model parameters θ as
∂LQ (ω,θ)ω ⋆ (θ) = arg minω LQ (ω, θ) =⇒= 0. Now we∂ωhave task loss which is optimized to find optimal Q-values:
Ltask (ω ⋆ (θ)) := Es,a [Qω⋆ (θ) (s, a) − BQω̄ (s, a)]2 (10)where the Bellman operator induced by ground-truth
trajectory andP reward is BQω̄ (s, a) := r(s, a) +γ Es,a,s′ log a′ exp Qω̄ (s′ , a′ ).

5. Conclusions
This paper addresses the challenge of combining task andprediction losses in task-based model learning. This 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 taskrelevant features and dimensions in the prediction space.

OMD setup. OMD end-to-end optimizes the model for thetask loss, i.e. θ⋆ = arg minθ Ltask (ω ⋆ (θ)).

TaskMet setup. For metric learning, we extend OMD tolearn a metric using task gradients, to train the model parameters, see Fig. 5. Metric learning just adds one more level of4

TaskMet: Task-Driven Metric Learning for Model Learning

References

Futoma, J., Hughes, M. C., and Doshi-Velez, F. Popcorn:
Partially observed prediction constrained reinforcement
learning. arXiv preprint arXiv:2001.04032, 2020. Cited
on page 3.

Agrawal, A., Amos, B., Barratt, S., Boyd, S., Diamond, S.,
and Kolter, J. Z. Differentiable convex optimization layers. Advances in neural information processing systems,
32, 2019. Cited on page 3.

Ghojogh, B., Ghodsi, A., Karray, F., and Crowley, M. Spectral, probabilistic, and deep metric learning: Tutorial andsurvey. arXiv preprint arXiv:2201.09267, 2022. Cited on
page 2.

Amos, B., Jimenez, I., Sacks, J., Boots, B., and Kolter, J. Z.
Differentiable mpc for end-to-end planning and control.
Advances in neural information processing systems, 31,
2018. Cited on page 1.

Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream tocontrol: Learning behaviors by latent imagination. arXiv
preprint arXiv:1912.01603, 2019a. Cited on page 2.

Bengio, Y. Using a financial training criterion rather thana prediction criterion. International journal of neuralsystems, 8(04):433–443, 1997. Cited on pages 1 and 2.
Bhardwaj, M., Boots, B., and Mukadam, M. Differentiable
gaussian process motion planning. In 2020 IEEE international conference on robotics and automation (ICRA), pp.
10598–10604. IEEE, 2020. Cited on page 1.

Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D.,
Lee, H., and Davidson, J. Learning latent dynamics forplanning from pixels. In International conference on
machine learning, pp. 2555–2565. PMLR, 2019b. Cited
on page 2.

Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S.,
Llinares-López, F., Pedregosa, F., and Vert, J.-P. Efficientand modular implicit differentiation. Advances in neural
information processing systems, 35:5230–5242, 2022.
Cited on page 3.

Hansen, N., Wang, X., and Su, H. Temporal difference
learning for model predictive control. arXiv preprint
arXiv:2203.04955, 2022. Cited on page 2.
Hastie, T. and Tibshirani, R. Discriminant adaptive nearest neighbor classification and regression. Advances inneural information processing systems, 8, 1995. Cited on
page 3.

Bottou, L., Curtis, F. E., and Nocedal, J. Optimization
methods for large-scale machine learning. SIAM review,60(2):223–311, 2018. Cited on page 2.
Dini, U. Analisi infinitesimale. Lithografia Gorani, 1878.
Cited on page 7.

Hauberg, S., Freifeld, O., and Black, M. A geometric takeon metric learning. Advances in Neural Information Processing Systems, 25, 2012. Cited on page 3.

Dontchev, A. L. and Rockafellar, R. T. Implicit functions
and solution mappings, volume 543. Springer, 2009.
Cited on page 7.

Kang, B., Liu, Z., Wang, X., Yu, F., Feng, J., and Darrell,
T. Few-shot object detection via feature reweighting. In
Proceedings of the IEEE/CVF International Conference
on Computer Vision, pp. 8420–8429, 2019. Cited on
pages 1 and 2.

Donti, P., Amos, B., and Kolter, J. Z. Task-based end-to-end
model learning in stochastic optimization. Advances inneural information processing systems, 30, 2017. Cited
on pages 1, 2, and 3.

Kaya, M. and Bilge, H. Ş. Deep metric learning: A survey.
Symmetry, 11(9):1066, 2019. Cited on page 3.

El Balghiti, O., Elmachtoub, A. N., Grigas, P., and Tewari,
A. Generalization bounds in the predict-then-optimize
framework. Advances in neural information processing
systems, 32, 2019. Cited on pages 1 and 2.

Kulis, B. et al. Metric learning: A survey. Foundations
and Trends® in Machine Learning, 5(4):287–364, 2013.
Cited on page 3.

Elmachtoub, A. N. and Grigas, P. Smart “predict, then
optimize”. Management Science, 68(1):9–26, 2022. Cited
on pages 1 and 2.

Liu, M., Grigas, P., Liu, H., and Shen, Z.-J. M. Active learning in the predict-then-optimize framework: A marginbased approach. arXiv preprint arXiv:2305.06584, 2023.
Cited on pages 1 and 2.

Farahmand, A.-m. Iterative value-aware model learning.
Advances in Neural Information Processing Systems, 31,
2018. Cited on page 1.

Mandi, J., Stuckey, P. J., Guns, T., et al. Smart predict-andoptimize for hard combinatorial optimization problems.
In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 1603–1610, 2020. Cited on
pages 1 and 2.

Farahmand, A.-m., Barreto, A., and Nikovski, D. Valueaware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics, pp. 1486–
1494. PMLR, 2017. Cited on pages 1 and 2.
5

TaskMet: Task-Driven Metric Learning for Model Learning

Markowitz, H. M. and Todd, G. P. Mean-variance analysis
in portfolio choice and capital markets, volume 66. John
Wiley & Sons, 2000. Cited on page 7.
Michaud, R. O. The markowitz optimization enigma: Is‘optimized’optimal? Financial analysts journal, 45(1):31–42, 1989. Cited on page 7.
Nikishin, E., Abachi, R., Agarwal, R., and Bacon, P.L. Control-oriented model-based reinforcement learning with implicit differentiation. In Proceedings of the
AAAI Conference on Artificial Intelligence, volume 36,
pp. 7886–7894, 2022. Cited on pages 1, 2, 4, 8, 9, and 10.
Sadana, U., Chenreddy, A., Delage, E., Forel, A., Frejinger,E., and Vidal, T. A survey of contextual optimization
methods for decision making under uncertainty, 2023.
Cited on page 2.
Shah, S., Wang, K., Wilder, B., Perrault, A., and Tambe,
M. Decision-focused learning without decision-making:Learning locally optimized decision losses. In Oh,
A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.),
Advances in Neural Information Processing Systems,
2022. URL https://openreview.net/forum?
id=eN2lQxjWL05. Cited on pages 1, 2, 3, 7, 8, and 9.
Voelcker, C., Liao, V., Garg, A., and Farahmand, A.m. Value gradient weighted model-based reinforcement
learning. arXiv preprint arXiv:2204.01464, 2022. Cited
on pages 1 and 2.
Weinberger, K. Q. and Tesauro, G. Metric learning forkernel regression. In Artificial intelligence and statistics,
pp. 612–619. PMLR, 2007. Cited on page 3.
Wilder, B., Dilkina, B., and Tambe, M. Melding the datadecisions pipeline: Decision-focused learning for combinatorial optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 1658–1665,
2019. Cited on pages 1, 2, and 3.
Xie, Y., Dai, H., Chen, M., Dai, B., Zhao, T., Zha, H.,
Wei, W., and Pfister, T. Differentiable top-k with optimaltransport. Advances in Neural Information Processing
Systems, 33:20520–20531, 2020. Cited on page 3.
Yang, L. and Jin, R. Distance metric learning: A comprehensive survey. Michigan State Universiy, 2(2):4, 2006.
Cited on page 3.

6

TaskMet: Task-Driven Metric Learning for Model Learning

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 & Rockafellar (2009, Theorem 1B.1):
Theorem 1 (Implicit Function Theorem). Implicit Function Theorem: Let f : Rn × Rm → R be a continuous differentiable⋆ ⋆,y )̸= 0, then there exists an open set
function, and let x⋆ , y ⋆ be a point satisfying f (x⋆ , y ⋆ ) = 0. If the Jacobian ∂(f (x∂y⋆ ⋆⋆around (x , y ) and a unique continuously differentiable function g such that y = g(x⋆ ) and f (x, g(x)) = 0. Additionally,
the following relation holds:


∂g(x)
∂f (x, y ⋆ ) ∂f (x, y ⋆ )=−|y⋆ =g(x)
(12)
∂x
∂y
∂x

B. Experimental Details
B.1. Metric Parameterization
We parameterize the metric using a neural network with parameters ϕ, denoted as Λϕ := L⊤ϕ Lϕ , where Lϕ is an n × nmatrix, where n is the dimension of the prediction space. This particular factorization constraint ensures that the metricmatrix is positive semi-definite, which is crucial for it to be considered a valid metric. To initialize the neural network
parameters, we set them to initialize the matrix closer to the identity matrix I, representing the Euclidean metric. The learned
metric can be conditional on the input, denoted as Λϕ (x), or unconditional, represented as Λϕ , depending on the problem’s
structure.
B.2. Decision Oriented Model Learning
Description of the resource allocation problems (Shah et al., 2022) used to benchmark our method:• Cubic This setting evaluate methods under model mismatch scenario where the model being learned suffer with severe
approximation error. In this task, it is important for methods to model the data well around the points that are morecritical for the downstream tasks. Prediction Model: A linear prediction model fθ (x) := θx is learned for the problem
3
where the ground truth data is generated by cubic function i.e. D := {xi , yi }N
i=1 yi = 10xi − 6.5xi , xi ∈ U [−1, 1].
⋆
Downstream task: Choose top B = 1 out of M = 50 resources ŷ = [ŷ1 , . . . , ŷM ], z (ŷ) := arg maxi ŷ• Budget Allocation: Choose top B = 2 websites to advertise based on Click-through-rates (CTRs) predictionsof K users on M websites. Prediction Model: ŷm = fθ (xm ) where xm is given features of mth website andŷm = [ŷm,1 , . . . , ŷm,K ] is the predicted CTRs for mth website for all K users. Downstream task: Determine
B = 2 websites such that expected number of users that click on the ad at least once is maximized i.e., z ⋆ (ŷm ) =PK
QM
arg maxz j=0 (1 − i=0 zi · ŷij ) where zi ∈ {0, 1}.

• Portfolio Optimization: The task is to choose a distribution over M stocks in Markowitz portfolio optimization(Markowitz & Todd, 2000; Michaud, 1989) that maximize the expected return under the risk penalty. Prediction Model:
Given the historical data xm about a stock m, predict the future stock price ŷm . Combining prediction over M stocks
to get ŷ = [ŷ1 , . . . , ŷM ]. Downstream Task: Given the correlation matrix Q of the stocks, choose a distribution overPMstocks z⋆ (ŷ) = arg maxz z⊤ ŷ − λz⊤ Qz s. t. i=0 zi ≤ 1 and 0 ≤ zi ≤ 1, ∀i
Description of different baselines considered for benchmarking:• MSE: Prediction model trained with standard MSE loss, θ⋆ = arg minθ Ex,y∼D [(fθ (x) − y)2 ]. This method doesn’t
use any task information.
• DFL: Prediction model trained with combination of Ltask and Lpred as shown in Eq. (1).
• LODL: Shah et al. (2022) learn parametric loss for each point in the training data to approximate the Ltask aroundthat point i.e., LODLψn (ŷn ) ≈ Ltask (ŷn )∀n. They create a dataset of {(ŷn , Ltask (ŷn ))} for ŷn sampled around the yn .
PK
1
k
k 2
After this they learn the LODL loss for each point as follows - ψn⋆ = arg minψn K
k=1 (LODLψn (yn ) − Ltask (yn )) .
⊤We chose the variant of their method which was closest to ours, where LODLψn (ŷ) = (ŷ − y) ψn (ŷ − y) where ψn
7

TaskMet: Task-Driven Metric Learning for Model Learning

is a learned symmetric Positive semidefinite (PSD) matrix. LODL also uses Eq. (1) to learn the model parameters, butusing LODLψn (ŷn ) ≈ Ltask (ŷn )Table 2. Prediction Error of different methods

Problems
MethodCubic

MSE
DFL (α = 0)
DFL (α = 10)
LODL-Quadratic (α = 0)
LODL-Quadratic( α = 10)
TaskMetBudget Allocation

2.30 ± 0.03
2.89 ± 0.32
2.41 ± 0.05
2.88 ± 0.030
2.29 ± 0.19
2.89 ± 0.03

−4

Portfolio Optimization−44.32e ± 2.35e
7.17e−3 ± 5.83e−3
8.09e−4 ± 12.07e−4
3.59e−3 ± 1.29e−3
5.05e−4 ± 1.88e−4
9.74e−4 ± 13.79e−4

4.03e−4 ± 0.24e−40.826 ± 0.081
5.18e−4 ± 0.46e−4
5.56e−3 ± 9.95e−4
4.31e−4 ± 0.31e−4
4.69e−4 ± 0.56e−4

B.3. Model-based reinforcement learning
Following is the derivation of final gradient to learn ϕ from Eq. (11). Using the implicit function theorem and using it on
Eq. (11), we get the following∂ω ⋆ ∂θ⋆·∂θ⋆ ∂ϕ−1∂L(ω, θ⋆ )∂L(ω, θ⋆ )=∇ω Ltask (ω ⋆ ) ·
·
2
∂ ω
∂θ∂ω

∇ϕ Ltask =∇ω Ltask (ω ⋆ ) ·(13)

ω=ω ⋆ (θ ⋆ )·∂Lpred (θ, ϕ)
∂2θ

−1

·

∂Lpred (θ, ϕ)
∂ϕ∂θ

(14)
θ=θ ⋆ (ϕ)

TaskMetLearning

MetricPlanning

Q*

Model
IFT

ActingLoss
Backprop

IFT
OMD

Figure 5. OMD (Nikishin et al., 2022) uses 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 tolearn model parameters.

We also consider a setting with reduced model capacity, where the network is under-parametrized, forcing the model toprioritize how it allocates its capacity. In this scenario, we employ a full conditional metric, denoted as Λϕ = Λϕ (x), whichenables the metric to weigh dimensions and state-action pairs differently. We conducted the experiment using a model sizeof 3 hidden units in the layer. As depicted in Fig. 6, 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 also contributes to a better dynamics model.

C. Implementation Details
C.1. Source code
Upon request, we will provide an anonymized version of our code in the rebuttal.
C.2. Decision Oriented Model Learning
We replicated our experiments using the codebase provided by Shah et al. (2022), which can be found at github. To ensure
consistency, we used the same hyperparameters as mentioned in the code or article for the baselines. Our metric learning
8

TaskMet: Task-Driven Metric Learning for Model Learning

500

Episode Return

Dynamics (MSE)

400
300
200

100

10 1

100
00

101

50K

100K

150K

Environment Steps

200K

TaskMet

0
MLE

50K

100K

150K

Environment Steps

200K

OMD

Figure 6. Results on cartpole with a reduced model capacity from Nikishin et al. (2022).

pipeline was added on top of their code, and thus we focused on tuning hyperparameters related to metric learning. The
metric is parameterized as Λϕ (x) = Lϕ (x)L⊤ϕ (x) + ϵϕ In×n , where ϵϕ is a learnable parameter that explicitly controls theamount of Euclidean metric in the predicted metric. This helps ensure the stability of metric learning. We initialize theparameters 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 withdifferent seeds for each of the experiments, where each method was evaluated on 10 different datasets, with 5 different seeds
used for each dataset.
Table 3. Hyper-parameters for Decision Oriented Learning Experiments

Hyper-ParameterΛϕ learning rate
Λϕ hidden layer sizes
Warmup steps
Inner Iterations (K)Implicit derivative batchsize
Implicit derivative solver

Values
10−3
[200]
500
100
10
Conjugate gradient on the normal equations (5 iterations)C.3. 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 sourcecode. 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 4. Hyper-parameters for the CartPole experiments

Hyper-ParameterΛϕ learning rate
Λϕ hidden layer sizes
Warmup steps
Inner iterations (K)Implicit derivative batchsize
Implicit derivative solver

Values10−3[32, 32]5000
1
256Conjugate gradient on the normal equations (10 iterations)

For the state distractor experiments, we parameterize the metric as an unconditional diagonal matrix, denoted as Λϕ =diag(ϕ) where ϕ ∈ Rn and n is the dimension of the state space. In addition, we also consider a hyper-parameter of9

TaskMet: Task-Driven Metric Learning for Model Learning

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 ϕ vector to be equal to the L2 norm of an euclidean metric
which isϕ √n,ensuring
used by MSE method. This constraint the family of learnable metrics. To achieve this, we set ϕ := ∥ϕ∥2√∥ϕ∥2 = ∥In×n ∥2 = 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 ofhyperparameters - [unnormalized, normalized] × [0.0, 0.001, 0.1] and choose the best hyper-parameter combination for eachof 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.

D. Training time analysis
We also measure the training time of our method compared to MLE-based training, as using implicit derivatives may
introduce additional computational overhead. It is important to note that the evaluation time is not affected by the use of
metric learning, as the metric is only employed during the training phase. Consequently, the evaluation runtime of the agent
remains the same for both methods.
Table 5. Time (sec) for one update step of the agent

MethodTime (s)MLE
TaskMet0.02400.0243We measure the time for single update step of agent. The code is written in JAX and evaluated on the computer with GPU— Quadro GV100 and CPU — Intel Xeon Gold 6230 @ 2.10GHz. As shown in Table 5, learning a metric using implicit
derivative to train the predictor takes negligible extra time compared to MLE method.

10