SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \operatorname {c o s t} \left(x _ {1: T}\right) \leq \left(1 + \lambda_ {2}\right) \operatorname {c o s t} \left(\tilde {x} _ {1: T}\right) + \frac {\beta + \alpha^ {2}}{2} \left(1 + \frac {1}{\lambda_ {2}}\right). \\ \sum_ {i = 1} ^ {T} \left(\left[ \| \tilde {x} _ {t} - x _ {t} ^ {\pi} \| - \sqrt {K \left(d \left(x _ {t} ^ {\pi} , x _ {t - p : t - 1} ^ {\pi}\right) + \sum_ {\tau \in \mathcal {D} _ {t}} f \left(x _ {\tau} ^ {\pi} , y _ {\tau}\right)\right)} \right] ^ {+}\right) ^ {2} \tag {45} \\ \end{array}
$$
By substituting  $K = \frac{2(\lambda - \lambda_0)}{(\beta_h + \alpha^2)(1 + \frac{1}{\lambda_0})}$  back to Eqn (46), we have
$$
\begin{array}{l} \operatorname {c o s t} \left(x _ {1: T}\right) \leq \left(1 + \lambda_ {2}\right) \operatorname {c o s t} \left(\tilde {x} _ {1: T}\right) + \left(1 + \frac {1}{\lambda_ {2}}\right) \sum_ {i = 1} ^ {T} \left[ \frac {\beta + \alpha^ {2}}{2} \| \tilde {x} _ {t} - x _ {t} ^ {\pi} \| ^ {2} \right. \tag {46} \\ \left. - \frac {\lambda - \lambda_ {0}}{1 + \frac {1}{\lambda_ {0}}} \left(d (x _ {t} ^ {\pi}, x _ {t - p: t - 1} ^ {\pi}) + \sum_ {\tau \in \mathcal {D} _ {t}} f (x _ {\tau} ^ {\pi}, y _ {\tau})\right) \right] ^ {+} \\ \end{array}
$$
By defining single step cost of the expert  $\pi$  as  $\mathrm{cost}_t^\pi = d(x_t^\pi, x_{t - p:t - 1}^\pi) + \sum_{\tau \in \mathcal{D}_t} f(x_\tau^\pi, y_\tau)$  and the auxiliary cost as  $\Delta(\lambda) = \sum_{i=1}^{T} \left[ \| \tilde{x}_t - x_t^\pi \|^2 - \frac{2(\lambda - \lambda_0)}{(\beta_h + \alpha^2)(1 + \frac{1}{\lambda_0})} \mathrm{cost}_t^\pi \right]^+$ .
$$
\operatorname {c o s t} \left(x _ {1: T}\right) \leq \left(\sqrt {\operatorname {c o s t} \left(\tilde {x} _ {1 : T}\right)} + \sqrt {\frac {\beta + \alpha^ {2}}{2} \Delta (\lambda)}\right) ^ {2} \tag {47}
$$
Combined with Lemma C.4, we obtain the following bound, which finished this proof.
$$
\operatorname {c o s t} \left(x _ {1: T}\right) \leq \min  \left((1 + \lambda) \operatorname {c o s t} \left(x _ {1: T} ^ {\pi}\right), \left(\sqrt {\operatorname {c o s t} \left(\tilde {x} _ {1 : T}\right)} + \sqrt {\frac {\beta + \alpha^ {2}}{2} \Delta (\lambda)}\right) ^ {2}\right) \tag {48}
$$
Proof. We first give the formal definition of Rademacher complexity of the ML model space with robustification.
Definition 5 (Rademacher Complexity). Let  $\operatorname{Rob}_{\lambda}(\mathcal{W}) = \{\operatorname{Rob}_{\lambda}(h_W), W \in \mathcal{W}\}$  be the ML model space with robustification constrained by (2). Given the dataset  $\mathcal{S}$ , the Rademacher complexity with respect to  $\operatorname{Rob}_{\lambda}(\mathcal{W})$  is
$$
\operatorname {R a d} _ {\mathcal {S}} (\operatorname {R o b} _ {\lambda} (\mathcal {W})) = \frac {1}{| \mathcal {S} |} \mathbb {E} _ {\nu} \left[ \sup  _ {W \in \mathcal {W}} \left(\sum_ {i \in \mathcal {S}} \nu_ {i} \operatorname {R o b} _ {\lambda} (h _ {W} (y ^ {i}))\right) \right],
$$
where  $y^{i}$  is the  $i$ -th sample in  $S$ , and  $\nu_{1},\dots,\nu_{n}$  are independently drawn from Rademacher distribution.
Since the cost functions are smooth, they are locally Lipschitz continuous for the bounded action space, and we can apply the generalization bound based on Rademacher complexity [60] for the space of robustified ML model  $\mathsf{Rob}_{\lambda}(h_W)$ . Given any ML model  $h_W$  trained on dataset  $S$ , with probability at least  $1 - \delta, \delta \in (0,1)$ ,
$$
\mathbb {E} _ {\mathbb {P} _ {y} ^ {\prime}} \left[ \operatorname {c o s t} _ {1: T} \right] \leq \overline {{\operatorname {c o s t}}} _ {\mathcal {S}} \left(\operatorname {R o b} _ {\lambda} \left(h _ {W}\right)\right) + 2 \Gamma_ {x} \operatorname {R a d} _ {\mathcal {S}} \left(\operatorname {R o b} _ {\lambda} (\mathcal {W})\right) + 3 \bar {c} \sqrt {\frac {\log (2 / \delta)}{| \mathcal {S} |}}, \tag {49}
$$
where  $\Gamma_{x} = \sqrt{T} |\mathcal{X}| \left[ \beta_{h} + \frac{1}{2} (1 + \sum_{i=1}^{p} L_{i}) (1 + \sum_{i=1}^{p} L_{i}) \right]$  with  $|\mathcal{X}|$  being the size of the action space  $\mathcal{X}$  and  $\beta_{h}, L_{i}$ , and  $p$  as the smoothness constant, Lipschitz constant of the nonlinear term in the switching cost, and the memory length as defined in Assumptions 1 and 2, and  $\bar{c}$  is the upper bound of the total cost for an episode. We can get the average cost bound in Proposition 4.2.
Next, we prove that the Rademacher complexity of the ML model space with robustification is no larger than the Rademacher complexity of the ML model space without robustification expressed as  $\{h_W, W \in \mathcal{W}\}$ , i.e. we need to prove  $\mathrm{Rad}_{\mathcal{S}}(\mathrm{Rob}_{\lambda}(\mathcal{W})) \leq \mathrm{Rad}_{\mathcal{S}}(\mathcal{W})$ . The Rademacher complexity can be expressed by Dudley's entropy integral [61] as
$$
\operatorname {R a d} _ {\mathcal {S}} \left(\operatorname {R o b} _ {\lambda} (\mathcal {W})\right) = \mathcal {O} \left(\frac {1}{\sqrt {| \mathcal {S} |}} \int_ {0} ^ {\infty} \sqrt {\log \mathbb {N} \left(\epsilon , \operatorname {R o b} _ {\lambda} (\mathcal {W}) , L _ {2} (\mathcal {S})\right)} \mathrm {d} \epsilon\right), \tag {50}
$$
where  $\mathbb{N}(\epsilon, \operatorname{Rob}_{\lambda}(\mathcal{W}), L_2(\mathcal{S}))$  is the covering number [61] with respect to radius  $\epsilon$  and the function distance metric  $\| h_1 - h_2 \|_{L_2(\mathcal{S})} = \frac{1}{|\mathcal{S}|} \sum_{i \in \mathcal{S}} \| h_1(x_i) - h_2(x_i) \|^2$  where  $h_1$  and  $h_2$  are two functions defined on the space including dataset  $\mathcal{S}$ . We can find that for any two different weights  $W_1$  and  $W_2$ , their corresponding post-robustification distance  $\| \operatorname{Rob}_{\lambda}(h_{W_1}) - \operatorname{Rob}_{\lambda}(h_{W_2}) \|_{L_2(\mathcal{S})}$  is no larger than their pre-robustification distance  $\| h_{W_1} - h_{W_2} \|_{L_2(\mathcal{S})}$ . To see this, we discuss three cases given any input sample  $y$ . If both  $h_{W_1}(y)$  and  $h_{W_2}(y)$  lie in the projection set, then  $\operatorname{Rob}_{\lambda}(h_{W_1})(y) = h_{W_1}(y)$  and  $\operatorname{Rob}_{\lambda}(h_{W_2})(y) = h_{W_2}(y)$ . If  $h_{W_1}(y)$  lies in the projection set while  $h_{W_2}(y)$  is out of the projection set, the projection operation based on the closed convex projection set will make  $\| \operatorname{Rob}_{\lambda}(h_{W_1})(y) - \operatorname{Rob}_{\lambda}(h_{W_2})(y) \|$  to be less than  $\| h_{W_1}(y) - h_{W_2}(y) \|$ . If both  $h_{W_1}(y)$  and  $h_{W_2}(y)$  lie out of the projection set, we still have  $\| \operatorname{Rob}_{\lambda}(h_{W_1})(y) - \operatorname{Rob}_{\lambda}(h_{W_2})(y) \| \leq \| h_{W_1}(y) - h_{W_2}(y) \|$ .
since the projection set at each round is a closed convex set [62]. Therefore, after robustification, the distance between two models with different weights will not become larger, i.e.  $\| \mathrm{Rob}_{\lambda}(h_{W_1}) - \mathrm{Rob}_{\lambda}(h_{W_2})\|_{L_2(\mathcal{S})}\leq \| h_{W_1} - h_{W_2}\|_{L_2(\mathcal{S})}$ , which means RCL has a covering number  $\mathbb{N}(\epsilon ,\operatorname {Rob}_{\lambda}(\mathcal{W}),L_2(\mathcal{S}))$  no larger than that of the individual ML model  $\mathbb{N}(\epsilon ,\mathcal{W},L_2(\mathcal{S}))$  for any  $\epsilon$ . Thus the Rademacher complexity with the robustification procedure does not increase.
By [63], the upper bound of Rademacher complexity with respect to the space of ML model  $\mathrm{Rad}_S(\mathrm{Rob}_\lambda (\mathcal{W}))$  is in the order of  $\mathcal{O}(\frac{1}{\sqrt{|S|}})$ . Since the Rademacher complexity with the robustification procedure satisfies  $\mathrm{Rad}_S(\mathrm{Rob}_\lambda (\mathcal{W})) \leq \mathrm{Rad}_S(\mathcal{W})$ , it also decreases with the dataset size in the order of  $\mathcal{O}(\frac{1}{\sqrt{|S|}})$ .
D Robustification-aware Training
Theorem 4.2 also shows the benefits of training the ML model in a robustification-aware manner. Specifically, by comparing the losses in (5) and (6), we see that using (6) as the robustification-aware loss for training  $W$  can reduce the term  $\overline{\mathrm{cost}}_S(\mathrm{ROB}(h_W))$  in the average cost bound, which matches exactly with the training objective in (6). The robustification-aware approach is only beginning to be explored in the ML-augmented algorithm literature and non-trivial (e.g., unconstrained downstream optimization in [55]), especially considering that (1) is a constrained optimization problem with no explicit gradients.
Gradient-based optimizers such as Adam [64] are the de facto state-of-the-art algorithms for training ML models, offering better optimization results, convergence, and stability compared to those non-gradient-based alternatives [65]. Thus, it is crucial to derive the gradients of the loss function with respect to the ML model weight  $W$  given the added robustification step.
Next, we derive the gradients of  $x_{t}$  with respect to  $\tilde{x}_{t}$ . For the convenience of presentation, we use the basic SOCO setting with a single-step switching cost and no hitting cost delay as an example, while noting that the same technique can be extended to derive gradients in more general settings. Specifically, for this setting, the pre-robustification prediction is given by  $\tilde{x}_{t} = h_{W}(\tilde{x}_{t - 1},y_{t})$ , where  $W$  denotes the ML model weight. Then, the actual post-robustification action  $x_{t}$  is obtained by projection in (1) by setting  $q = 0$  and  $p = 1$ , given the ML prediction  $\tilde{x}_{t}$ , the expert's action  $x_{t}^{\pi}$  and cumulative cost  $(x_{1:t}^{\pi})$  up to  $t$ , and the actual cumulative cost  $(x_{1:t - 1})$  up to  $t - 1$ .
The gradient of the loss function  $\mathrm{cost}(x_{1:T}) = \sum_{t=1}^{T} (f(x_t, y_t) + d(x_t, x_{t-1}))$  with respect to the ML model weight  $W$  is given by  $\sum_{t=1}^{T} \nabla_W(f(x_t, y_t) + d(x_t, x_{t-1}))$ . Next, we write the gradient of per-step cost with respect to  $W$  as follows:
$$
\begin{array}{l} \nabla_ {W} \left(f (x _ {t}, y _ {t}) + d (x _ {t}, x _ {t - 1})\right) \\ = \nabla_ {x _ {t}} \left(f \left(x _ {t}, y _ {t}\right) + d \left(x _ {t}, x _ {t - 1}\right)\right) \nabla_ {W} x _ {t} + \nabla_ {x _ {t - 1}} \left(f \left(x _ {t}, y _ {t}\right) + d \left(x _ {t}, x _ {t - 1}\right)\right) \nabla_ {W} x _ {t - 1} \tag {51} \\ = \nabla_ {x _ {t}} \left(f (x _ {t}, y _ {t}) + d (x _ {t}, x _ {t - 1})\right) \nabla_ {W} x _ {t} + \nabla_ {x _ {t - 1}} d (x _ {t}, x _ {t - 1}) \nabla_ {W} x _ {t - 1}, \\ \end{array}
$$
where the gradients  $\nabla_{x_t}\big(f(x_t,y_t) + d(x_t,x_{t - 1})\big)$  and  $\nabla_{x_{t - 1}}d(x_t,x_{t - 1})$  are trivial given the hitting and switching cost functions, and the gradient  $\nabla_W x_{t - 1}$  is obtained at time  $t - 1$  in the same way as  $\nabla_W x_t$ . To derive  $\nabla_W x_t$ , by the chain rule, we have:
$$
\nabla_ {W} x _ {t} = \nabla_ {\tilde {x} _ {t}} x _ {t} \nabla_ {W} \tilde {x} _ {t} + \nabla_ {\operatorname {c o s t} (x _ {1: t - 1})} x _ {t} \nabla_ {W} \operatorname {c o s t} (x _ {1: t - 1}), \tag {52}
$$
where  $\nabla_W\tilde{x}_t$  is the gradient of the ML output (following a recurrent architecture illustrated in Fig. 1 in the appendix) with respect to the weight  $W$  and can be obtained recursively by using off-the-shelf BPTT optimizers [64], and  $\nabla_W\mathrm{cost}(x_{1:t - 1}) = \sum_{\tau = 1}^{t - 1}\nabla_W\big(f(x_\tau ,y_\tau) + d(x_\tau ,x_{\tau -1})\big)$  can also be recursively calculated once we have the gradient in Eqn. (51). Nonetheless, it is non-trivial to calculate the two gradient terms in Eqn. (52), i.e.,  $\nabla_{\tilde{x}_t}x_t$  and  $\nabla_{\mathrm{cost}(x_{1:t - 1})}x_t$ , where  $x_{t}$  itself is the solution to the constrained optimization problem (1) unlike in the simpler unconstrained case [55]. As we cannot explicitly write  $x_{t}$  in a closed form in terms of  $\tilde{x}_t$  and  $\mathrm{cost}(x_{1:t - 1})$ , we leverage the KKT conditions [66, 67, 68] to implicitly derive  $\nabla_{\tilde{x}_t}x_t$  and  $\nabla_{\mathrm{cost}(x_{1:t - 1})}x_t$  in the next proposition.
Proposition D.1 (Gradients by KKT conditions). Let  $x_{t} \in \mathcal{X}$  and  $\mu \geq 0$  be the primal and dual solutions to the problem (1), respectively. The gradients of  $x_{t}$  with respect to  $\tilde{x}_{t}$  and  $\mathrm{cost}(x_{1:t - 1})$  are
$$
\nabla_ {\bar {x} _ {t}} x _ {t} = \Delta_ {1 1} ^ {- 1} [ I + \Delta_ {1 2} S c (\Delta , \Delta_ {1 1}) ^ {- 1} \Delta_ {2 1} \Delta_ {1 1} ^ {- 1} ],
$$
$$
\nabla_ {c o s t (x _ {1: t - 1})} x _ {t} = \Delta_ {1 1} ^ {- 1} \Delta_ {1 2} S c (\Delta , \Delta_ {1 1}) ^ {- 1} \mu ,
$$
where  $\Delta_{11} = I + \mu \left( \nabla_{x_t, x_t} f(x_t, y_t) + \left( 1 + (1 + \frac{1}{\lambda_0}) (L_1^2 + L_1) \right) I \right)$ ,  $\Delta_{12} = \nabla_{x_t} f(x_t, y_t) + (x_t - \delta(x_{t-1})) + \left( 1 + (1 + \frac{1}{\lambda_0}) (L_1^2 + L_1) \right) (x_t - x_t^\pi)$ ,  $\Delta_{21} = \mu \Delta_{12}^\top$ ,  $\Delta_{22} = f(x_t, y_t) + d(x_t, x_{t-1}) + G(x_t, x_t^\pi) + cost(x_{1:t-1}) - (1 + \lambda) cost(x_{1:t}^\pi)$ , and  $Sc(\Delta, \Delta_{11}) = \Delta_{22} - \Delta_{21} \Delta_{11}^{-1} \Delta_{12}$  is the Schur-complement of  $\Delta_{11}$  in the blocked matrix  $\Delta = [[\Delta_{11}, \Delta_{12}], [\Delta_{21}, \Delta_{22}]]$ .
If the ML prediction  $\tilde{x}_t$  happens to lie on the boundary such that the inequality in (1) becomes an equality for  $x = \tilde{x}_t$ , then the gradient does not exist in this case and  $Sc(\Delta, \Delta_{11})$  may not be full-rank. Nonetheless, we can still calculate the pseudo-inverse of  $Sc(\Delta, \Delta_{11})$  and use Proposition D.1 to calculate the subgradient. Such approximation is actually a common practice to address non-differentiable points for training ML models, e.g., using 0 as the subgradient of  $ReLU(\cdot)$  at the zero point [64].