SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiao Zhang $^{1,2}$ , Ninglu Shao $^{1,2,*}$ , Zihua Si $^{1,2,*}$ , Jun Xu $^{1,2,\dagger}$ , Wenhan Wang $^{3}$ , Hanjing Su $^{3}$ , Ji-Rong Wen $^{1,2}$
$^{1}$  Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China  
 $^{2}$  Beijing Key Laboratory of Big Data Management and Analysis Methods, Beijing, China  
 $^{3}$  Tencent Inc., Shenzhen, China
{zhangx89, ninglu_shao, zihua_si, junxu, jrwen}@ruc.edu.cn {justinsu, ezewang}@tencent.com
Contextual batched bandit (CBB) is a setting where a batch of rewards is observed from the environment at the end of each episode, but the rewards of the non-executed actions are unobserved, resulting in partial-information feedback. Existing approaches for CBB often ignore the rewards of the non-executed actions, leading to underutilization of feedback information. In this paper, we propose an efficient approach called Sketched Policy Updating with Imputed Rewards (SPUIR) that completes the unobserved rewards using sketching, which approximates the full-information feedbacks. We formulate reward imputation as an imputation regularized ridge regression problem that captures the feedback mechanisms of both executed and non-executed actions. To reduce time complexity, we solve the regression problem using randomized sketching. We prove that our approach achieves an instantaneous regret with controllable bias and smaller variance than approaches without reward imputation. Furthermore, our approach enjoys a sublinear regret bound against the optimal policy. We also present two extensions, a rate-scheduled version and a version for nonlinear rewards, making our approach more practical. Experimental results show that SPUIR outperforms state-of-the-art baselines on synthetic, public benchmark, and real-world datasets.
Contextual bandits have gained significant popularity in solving sequential decision-making problems (Li et al., 2010; Lan and Baraniuk, 2016; Yom-Tov et al., 2017; Yang et al., 2021), where the agent continuously updates its decision-making policy fully online (i.e., at each step), considering the context and the received reward feedback to maximize cumulative rewards. In this paper, we address a more general setting called contextual batched bandits (CBB). In CBB, the decision process is divided into  $N$  episodes, and within each episode, the agent interacts with the environment for a fixed number of  $B$  steps. At the end of each episode, the agent collects reward feedbacks and contexts. Subsequently, the policy is updated using the collected data to guide the decision-making process in the subsequent episode. CBB offers a practical framework for real-world streaming applications (e.g., streaming recommendation (Zhang et al., 2021, 2022)). In the context of CBB settings, the batch size  $B$ , can be adjusted by the agent to achieve improved regret guarantees and meet the data throughput requirements based on the available computing resources (Zhou, 2023).
* Ninglu Shao and Zihua Si have made equal contributions to this paper.
† Corresponding author: Jun Xu.
In bandit settings, it is common for the environment to only provide feedback on the rewards of executed actions to the agent, while concealing the rewards of non-executed actions. This type of limited feedback is referred to as partial-information feedback (also called "bandit feedback"). In CBB setting, existing approaches tend to overlook the potential rewards associated with non-executed actions. Instead, they address the challenge of partial-information feedback through an exploration-exploitation tradeoff in both the context space and reward space (Han et al., 2020; Zhang et al., 2020). However, CBB agents typically estimate and maintain reward models for the action-selection policy, thereby capturing some information about the potential rewards of non-executed actions. This additional reward structure information is available for policy updating in each episode but remains untapped by existing batched bandit approaches.
In the context of contextual bandit settings where the policy is updated online, several bias-correction approaches have been introduced to tackle the issue of partial-information feedback. Dimakopoulou et al. (2019) presented linear contextual bandits integrating the balancing approach from causal inference, which reweight the contexts and rewards by the inverse propensity scores. Chou et al. (2015) designed pseudo-reward algorithms for contextual bandits using upper confidence bound (UCB) strategy, which use a direct method to estimate the unobserved rewards. Kim and Paik (2019) focused on the correction of feedback bias for LASSO bandit with high-dimensional contexts, and applied the doubly-robust approach to the reward modification using average contexts. While these approaches have demonstrated effectiveness in contextual bandit settings, little attention has been given to addressing the under-utilization of partial-information feedback in CBB setting.
Theoretical and experimental analyses in Section 2 indicate that better performance of CBB is achievable if the rewards of the non-executed actions can be received. Motivated by these observations, we propose a novel reward imputation approach for the non-executed actions, which mimics the reward generation mechanisms of environments. We conclude our contributions as follows.
(1) To fully utilize feedback information in CBB, we formulate the reward imputation as a problem of imputation regularized ridge regression, where the policy can be updated efficiently using sketching.
(2) We prove that our reward imputation approach obtains a relative-error bound for sketching approximation, achieves an instantaneous regret with a controllable bias and a smaller variance than that without reward imputation, has a lower bound of the sketch size independently of the overall number of steps, enjoys a sublinear regret bound against the optimal policy, and reduces the time complexity from  $O(Bd^2)$  to  $O(cd^2)$  for each action in one episode, where  $B$  denotes the batch size,  $c$  the sketch size, and  $d$  the dimension of the context space, satisfying  $d < c < B$ .
(3) We present two practical variants of our reward imputation approach, including the rate-scheduled version that sets the imputation rate without tuning, and the version for nonlinear rewards.
(4) We carried out extensive experiments on a synthetic dataset, the publicly available Criteo dataset, and a dataset from a commercial app to demonstrate our performance, empirically analyzed the influence of different parameters, and verified the correctness of the theoretical results.
Related Work. Recently, batched bandit has become an active research topic in statistics and learning theory including 2-armed bandit (Perchet et al., 2016), multi-armed bandit (Gao et al., 2019; Zhang et al., 2020; Wang and Cheng, 2020), and contextual bandit (Han et al., 2020; Ren and Zhou, 2020; Gu et al., 2021). Han et al. (2020) defined linear contextual bandits, and designed UCB-type algorithms for both stochastic and adversarial contexts, where true rewards of different actions have the same parameters. Zhang et al. (2020) provided methods for inference on data collected in batches using bandits, and introduced a batched least squares estimator for both multi-arm and contextual bandits. Recently, Esfandiari et al. (2021) proved refined regret upper bounds of batched bandits in stochastic and adversarial settings. There are several recent works that consider similar settings to CBB, e.g., episodic Markov decision process (Jin et al., 2018), LASSO bandits (Wang and Cheng, 2020). Sketching is another related technology that compresses a large matrix to a much smaller one by multiplying a (usually) random matrix while retaining certain properties (Woodruff, 2014), which has been used in online convex optimization (Calandriello et al., 2017; Zhang and Liao, 2019).
2 Problem Formulation and Analysis
Let  $[x] = \{1,2,\dots ,x\}$ ,  $S\subseteq \mathbb{R}^d$  be the context space whose dimension is  $d$ ,  $\mathcal{A} = \{A_j\}_{j\in [M]}$  the action space containing  $M$  actions,  $[A;B] = [A^{\intercal},B^{\intercal}]^{\intercal}$ ,  $\| A\|_{\mathrm{F}}$ ,  $\| A\| _1\| A\| _2$  denote the Frobenius
Protocol 1 Contextual Batched Bandit (CBB)
INPUT: Batch size  $B$  , number of episodes  $N$  , action space  $\mathcal{A} = \{A_j\}_{j\in [M]}$  , context space  $S\subseteq \mathbb{R}^d$    
1: Initialize policy  $p_0\gets 1 / M$  , sample data buffer  $\mathcal{D}_1 = \{(s_{0,b},A_{I_0,b},R_{0,b})\}_{b\in [B]}$  using initial policy  $p_0$    
2: for  $n = 1$  to  $N$  do   
3: Update the policy  $p_n$  on  $\mathcal{D}_n$    
4: for  $b = 1$  to  $B$  do   
5: Observe context  $s_{n,b}$  and choose  $A_{I_{n,b}}\in \mathcal{A}$  following the updated policy  $p_n(\boldsymbol{s}_{n,b})$    
6: end for   
7:  $\mathcal{D}_{n + 1}\leftarrow \{(s_{n,b},A_{I_{n,b}},R_{n,b})\}_{b\in [B]}$  , where  $R_{n,b}$  denotes the reward of action  $A_{I_{n,b}}$  on context  $s_{n,b}$    
8: end for
norm, 1-norm, and spectral norm of a matrix  $\mathbf{A}$ , respectively,  $\| \mathbf{a}\| _1$  and  $\| \mathbf{a}\| _2$  be the  $\ell_1$ -norm and the  $\ell_2$ -norm of a vector  $\mathbf{a}$ ,  $\sigma_{\mathrm{min}}(\mathbf{A})$  and  $\sigma_{\mathrm{max}}(\mathbf{A})$  denote the minimum and maximum of the of singular values of  $\mathbf{A}$ . In this paper, we focus on the setting of Contextual Batched Bandits (CBB) in Protocol 1, where the decision process is partitioned into  $N$  episodes, and in each episode, CBB consists of two phases: (1) the policy updating approximates the optimal policy based on the received contexts and rewards; (2) the online decision chooses actions for execution following the updated and fixed policy  $p$  for  $B$  steps ( $B$  is also called the batch size), and stores the context-action pairs and the observed rewards of the executed actions into a data buffer  $\mathcal{D}$ . The reward  $R$  in CBB is a partial-information feedback where rewards are unobserved for the non-executed actions.
In contrast to the existing batched bandit setting (Han et al., 2020; Esfandiari et al., 2021), where the true reward feedbacks for all actions are controlled by the same parameter vector while the received contexts differ across actions at each step, we make the assumption that in CBB setting, the mechanism of true reward feedback varies across actions, while the received context is shared among actions. Formally, for any context  $s_i \in S \subseteq \mathbb{R}^d$  and action  $A \in \mathcal{A}$ , we assume that the expectation of the true reward  $R_{i,A}^{\mathrm{true}}$  is determined by an unknown action-specific reward parameter vector  $\pmb{\theta}_{A}^{*} \in \mathbb{R}^{d}$ :  $\mathbb{E}[R_{i,A}^{\mathrm{true}} \mid s_i] = \langle \pmb{\theta}_A^*, s_i \rangle$  (the linear reward will be extended to the nonlinear case in Section 5). This setting for reward feedback matches many real-world applications, e.g., each action corresponds to a different category of candidate coupons in coupon recommendation, and the reward feedback mechanism of each category differs due to the different discount pricing strategies.
Next, we delve deeper into understanding the impact of unobserved feedbacks on the performance of policy updating in CBB setting. We first conducted an empirical comparison by applying the batch UCB policy (SBUCB) (Han et al., 2020) to environments under different proportions of received reward feedbacks. In particular, the agent under full-information feedback can receive all the rewards of the executed and non-executed actions, called Full-Information  $CB-B$  (FI-CBB) setting. From Figure 1, we can observe that the partial-information feedbacks are damaging in terms of hurting the policy updating, and batched bandit policy can benefit from more reward feedbacks, where the performance of  $80\%$  feedback is very close to that of FI-CBB.
Then, we prove the difference of instantaneous regrets between the CBB and FI-CBB settings in Theorem 1 (proof can be found in Appendix A).
Theorem 1. For any action  $A \in \mathcal{A}$  and context  $\boldsymbol{s}_i \in S$ , let  $\theta_A^n$  be the reward parameter vector estimated by the batched UCB policy in the  $n$ -th episode. The upper bound of instantaneous regret (defined by  $|\langle \theta_A^n, \boldsymbol{s}_i \rangle - \langle \theta_A^*, \boldsymbol{s}_i \rangle|$ ) in the FI-CBB setting is tighter than that in CBB setting (i.e., using the partial-information feedback).
From Theorem 1, we can infer that utilizing partial-information feedbacks leads to a deterioration in the regret of the bandit policy. Ideally, the policy would be updated using full-information feedback. However, in CBB, full-information feedback is unavailable. Fortunately, in CBB, different reward parameter vectors are maintained and estimated separately for each action, and the potential reward
Figure 1: Average rewards of batch UCB policy (Han et al., 2020) under different proportions of received reward feedbacks, interacting with the synthetic environment in Section 6, where  $x\%$  feedback means that  $x\%$  of actions can receive their true rewards
Figure 2: Comparison of the stored data corresponding to the action  $A_{j} \in \mathcal{A} = \{A_{j}\}_{j \in [M]}$  in CBB, CBB with our reward imputation, and full-information CBB, in the  $(n + 1)$ -th episode
structures of the non-executed actions have been captured to some extent. Therefore, why not utilize these maintained reward parameters to estimate the unknown rewards for the non-executed actions? In the following, we propose an efficient reward imputation approach that leverages this additional reward structure information to enhance the performance of the bandit policy.
3 Reward Imputation for Policy Updating
In this section, we present an efficient reward imputation approach tailored for policy updating in CBB setting.
Formulation of Reward Imputation. As shown in Figure 2, in contrast to CBB that ignores the contexts and rewards of the non-executed steps of each action, our reward imputation approach completes the missing values using the imputed contexts and rewards, approximating the full-information CBB setting. Specifically, at the end of the  $(n + 1)$ -th episode, for each action  $A_{j} \in \mathcal{A}, j \in [M]$ , the context vectors and rewards received at the steps where the action  $A_{j}$  is executed are observed, and are stored into a context matrix  $\pmb{S}_{A_{j}}^{n} \in \mathbb{R}^{N_{j}^{n} \times d}$  and a reward vector  $\pmb{R}_{A_{j}}^{n} \in \mathbb{R}^{N_{j}^{n}}$ , respectively, where  $N_{j}^{n}$  denotes the number of executed steps of  $A_{j}$  in episode  $n + 1$ . More importantly, at the steps (in episode  $n + 1$ ) where the action  $A_{j}$  is NOT executed, the following imputations need to be performed for action  $A_{j}$ : (1) since the contexts are shared by all the actions, we directly store them into an imputed context matrix  $\widehat{\pmb{S}}_{A_{j}}^{n} \in \mathbb{R}^{\widehat{N}_{j}^{n} \times d}$ , where  $\widehat{N}_{j}^{n}$  denotes the number of non-executed steps of  $A_{j}$  (i.e.,  $\widehat{N}_{j}^{n} = B - N_{j}^{n}$ ); (2) since the rewards of  $A_{j}$  are unobserved at the non-executed steps, we estimate them using an imputed reward vector: for any  $j \in [M]$ ,
$$
\widehat {\boldsymbol {R}} _ {A _ {j}} ^ {n} = \{r _ {n, 1} (A _ {j}), r _ {n, 2} (A _ {j}), \ldots , r _ {n, \widehat {N} _ {j} ^ {n}} (A _ {j}) \} \in \mathbb {R} ^ {\widehat {N} _ {j} ^ {n}},
$$
where  $r_{n,b}(A_j) \coloneqq \langle \bar{\theta}_{A_j}^n, s_{n,b} \rangle$  denotes the imputed reward parameterized by  $\bar{\theta}_{A_j}^n \in \mathbb{R}^d$  and  $s_{n,b}$  is the  $b$ -th row of  $\widehat{S}_{A_j}^n$ .
Next, we introduce the updating process of the reward parameter vector  $\bar{\theta}_{A_j}^n$ . We first concatenate the context and reward matrices from the previous episodes:  $L_{A_j}^n = [S_{A_j}^0;\dots ;S_{A_j}^n ]\in$ $\mathbb{R}^{L_j^n\times d}$ ,  $T_{A_j}^n = [R_{A_j}^0;\dots ;R_{A_j}^n ]\in \mathbb{R}^{L_j^n}$ ,  $L_{j}^{n} = \sum_{k = 0}^{n}N_{j}^{k}$ ,  $\widehat{L}_{A_j}^n = [\widehat{S}_{A_j}^0;\dots ;\widehat{S}_{A_j}^n ]\in$ $\mathbb{R}^{\widehat{L}_j^n\times d}$ ,  $\widehat{T}_{A_j}^n = [\widehat{R}_{A_j}^0;\dots ;\widehat{R}_{A_j}^n ]\in \mathbb{R}^{\widehat{L}_j^n}$ ,  $\widehat{L}_j^n = \sum_{k = 0}^{n}\widehat{N}_j^k$ . Then, the updated parameter vector  $\bar{\theta}_{A_j}^{n + 1}$  of the imputed reward for action  $A_{j}$  can be obtained by solving the following imputation regularized ridge regression: for  $n = 0,1,\ldots ,N - 1$
$$
\bar {\boldsymbol {\theta}} _ {A _ {j}} ^ {n + 1} = \underset {\boldsymbol {\theta} \in \mathbb {R} ^ {d}} {\arg \min } \underbrace {\left\| \boldsymbol {L} _ {A _ {j}} ^ {n} \boldsymbol {\theta} - \boldsymbol {T} _ {A _ {j}} ^ {n} \right\| _ {2} ^ {2}} _ {\text {O b s e r v e d T e r m}} + \underbrace {\gamma \left\| \widehat {\boldsymbol {L}} _ {A _ {j}} ^ {n} \boldsymbol {\theta} - \widehat {\boldsymbol {T}} _ {A _ {j}} ^ {n} \right\| _ {2} ^ {2}} _ {\text {I m p u t a t i o n T e r m}} + \lambda \| \boldsymbol {\theta} \| _ {2} ^ {2}, \tag {1}
$$
where  $\gamma \in [0,1]$  is the imputation rate that controls the degree of reward imputation and measures a trade-off between bias and variance (Remark 1&2),  $\lambda > 0$  is the regularization parameter. The discounted variant of the closed least squares solution of (1) is used for computing  $\bar{\theta}_{A_j}^{n+1}$ :
$$
\bar {\boldsymbol {\theta}} _ {A _ {j}} ^ {n + 1} = \left(\boldsymbol {\Psi} _ {A _ {j}} ^ {n + 1}\right) ^ {- 1} \left(\boldsymbol {b} _ {A _ {j}} ^ {n + 1} + \gamma \hat {\boldsymbol {b}} _ {A _ {j}} ^ {n + 1}\right), \tag {2}
$$
where  $\Psi_{A_j}^{n + 1}\coloneqq \lambda I_d + \Phi_{A_j}^{n + 1} + \gamma \widehat{\Phi}_{A_j}^{n + 1}$ , and
$$
\boldsymbol {\Phi} _ {A _ {j}} ^ {n + 1} = \boldsymbol {\Phi} _ {A _ {j}} ^ {n} + \boldsymbol {S} _ {A _ {j}} ^ {n \intercal} \boldsymbol {S} _ {A _ {j}} ^ {n}, \boldsymbol {b} _ {A _ {j}} ^ {n + 1} = \boldsymbol {b} _ {A _ {j}} ^ {n} + \boldsymbol {S} _ {A _ {j}} ^ {n \intercal} \boldsymbol {R} _ {A _ {j}} ^ {n}, \tag {3}
$$
$$
\widehat {\boldsymbol {\Phi}} _ {A _ {j}} ^ {n + 1} = \eta \widehat {\boldsymbol {\Phi}} _ {A _ {j}} ^ {n} + \widehat {\boldsymbol {S}} _ {A _ {j}} ^ {n \top} \widehat {\boldsymbol {S}} _ {A _ {j}} ^ {n}, \widehat {\boldsymbol {b}} _ {A _ {j}} ^ {n + 1} = \eta \widehat {\boldsymbol {b}} _ {A _ {j}} ^ {n} + \widehat {\boldsymbol {S}} _ {A _ {j}} ^ {n \top} \widehat {\boldsymbol {R}} _ {A _ {j}} ^ {n}, \tag {4}
$$
and  $\eta \in (0,1)$  is the discount parameter that controls how fast the previous imputed rewards are forgotten, and can help guaranteeing the regret bound in Theorem 2.
Efficient Reward Imputation using Sketching. As shown in the first 4 columns in Table 1, the overall time complexity of the imputation for each action is  $O(Bd^2)$  in each episode, where  $B$  represents the batch size, and  $d$  the dimensionality of the input. Thus, for all the  $M$  actions in one episode, reward imputation increases the time complexity from  $O(Bd^2)$  of the approach without imputation to  $O(MBd^2)$ . To address this issue, we design an efficient reward imputation approach using sketching, which reduces the time complexity of each action in one episode from  $O(Bd^2)$  to  $O(cd^2)$ , where  $c$  denotes the sketch size satisfying  $d < c < B$  and  $cd > B$ . Specifically, in the  $(n + 1)$ -th episode, the formulation in (1) can be approximated by a sketched ridge regression as:
$$
\tilde {\boldsymbol {\theta}} _ {A _ {j}} ^ {n + 1} = \underset {\boldsymbol {\theta} \in \mathbb {R} ^ {d}} {\arg \min } \left\| \boldsymbol {\Pi} _ {A _ {j}} ^ {n} \left(\boldsymbol {L} _ {A _ {j}} ^ {n} \boldsymbol {\theta} - \boldsymbol {T} _ {A _ {j}} ^ {n}\right) \right\| _ {2} ^ {2} + \gamma \left\| \widehat {\boldsymbol {\Pi}} _ {A _ {j}} ^ {n} \left(\widehat {\boldsymbol {L}} _ {A _ {j}} ^ {n} \boldsymbol {\theta} - \widehat {\boldsymbol {T}} _ {A _ {j}} ^ {n}\right) \right\| _ {2} ^ {2} + \lambda \| \boldsymbol {\theta} \| _ {2} ^ {2}, \tag {5}
$$
where  $\tilde{\theta}_A^{n + 1}$  denotes the updated parameter vector of the imputed reward using sketching for action  $A\in \mathcal{A}$ ,  $C_{A_j}^n\in \mathbb{R}^{c\times N_j^n}$  and  $\widehat{C}_{A_j}^n\in \mathbb{R}^{c\times \widehat{N}_j^n}$  are the sketch submatrices for the observed term and the imputation term, respectively, and the sketch matrices for the two terms can be represented as
$$
\boldsymbol {\Pi} _ {A _ {j}} ^ {n} = \left[ \boldsymbol {C} _ {A _ {j}} ^ {0}, \boldsymbol {C} _ {A _ {j}} ^ {1}, \dots , \boldsymbol {C} _ {A _ {j}} ^ {n} \right] \in \mathbb {R} ^ {c \times L _ {j} ^ {n}}, \quad \widehat {\boldsymbol {\Pi}} _ {A _ {j}} ^ {n} = \left[ \widehat {\boldsymbol {C}} _ {A _ {j}} ^ {0}, \widehat {\boldsymbol {C}} _ {A _ {j}} ^ {1}, \dots , \widehat {\boldsymbol {C}} _ {A _ {j}} ^ {n} \right] \in \mathbb {R} ^ {c \times \widehat {L} _ {j} ^ {n}}.
$$
We denote the sketches of the context matrix and the reward vector by  $\Gamma_{A_j}^n \coloneqq C_{A_j}^n S_{A_j}^n \in \mathbb{R}^{c \times d}$  and  $\Lambda_{A_j}^n \coloneqq C_{A_j}^n R_{A_j}^n \in \mathbb{R}^c$ , the sketches of the imputed context matrix and the imputed reward vector by  $\widehat{\Gamma}_{A_j}^n \coloneqq \widehat{C}_{A_j}^n \widehat{S}_{A_j}^n \in \mathbb{R}^{c \times d}$  and  $\widehat{\Lambda}_{A_j}^n \coloneqq \widehat{C}_{A_j}^n \widehat{R}_{A_j}^n \in \mathbb{R}^c$ . Similarly to (2), the discounted variant of the closed solution of (5) as follows:
$$
\tilde {\boldsymbol {\theta}} _ {A _ {j}} ^ {n + 1} = \left(\boldsymbol {W} _ {A _ {j}} ^ {n + 1}\right) ^ {- 1} \left(\boldsymbol {p} _ {A _ {j}} ^ {n + 1} + \gamma \hat {\boldsymbol {p}} _ {A _ {j}} ^ {n + 1}\right), \tag {6}
$$
where  $\eta \in (0,1)$  denotes the discount parameter,  $W_{A_j}^{n + 1}\coloneqq \lambda I_d + G_{A_j}^{n + 1} + \gamma \widehat{G}_{A_j}^{n + 1}$ , and
$$
\boldsymbol {G} _ {A _ {j}} ^ {n + 1} = \boldsymbol {G} _ {A _ {j}} ^ {n} + \boldsymbol {\Gamma} _ {A _ {j}} ^ {n \intercal} \boldsymbol {\Gamma} _ {A _ {j}} ^ {n}, \boldsymbol {p} _ {A _ {j}} ^ {n + 1} = \boldsymbol {p} _ {A _ {j}} ^ {n} + \boldsymbol {\Gamma} _ {A _ {j}} ^ {n \intercal} \boldsymbol {\Lambda} _ {A _ {j}} ^ {n}, \tag {7}
$$
$$
\widehat {\boldsymbol {G}} _ {A _ {j}} ^ {n + 1} = \eta \widehat {\boldsymbol {G}} _ {A _ {j}} ^ {n} + \widehat {\boldsymbol {\Gamma}} _ {A _ {j}} ^ {n \intercal} \widehat {\boldsymbol {\Gamma}} _ {A _ {j}} ^ {n}, \widehat {\boldsymbol {p}} _ {A _ {j}} ^ {n + 1} = \eta \widehat {\boldsymbol {p}} _ {A _ {j}} ^ {n} + \widehat {\boldsymbol {\Gamma}} _ {A _ {j}} ^ {n \intercal} \widehat {\boldsymbol {\Lambda}} _ {A _ {j}} ^ {n}. \tag {8}
$$
Using the parameter  $\tilde{\pmb{\theta}}_{A_j}^{n + 1}$ , we obtain the sketched version of imputed reward as  $\tilde{r}_{n,b}(A_j) := \langle \tilde{\pmb{\theta}}_{A_j}^n, \pmb{s}_{n,b} \rangle$  at step  $b \in [\widehat{N}_j^n]$ . Finally, we specify that the sketch submatrices  $\{C_A^n\}_{A \in \mathcal{A}, n \in [N]}$  and  $\{\widehat{C}_{A_j}^n\}_{A \in \mathcal{A}, n \in [N]}$  are the block construction of Sparser Johnson-Lindenstrauss Transform (SJLT) (Kane and Nelson, 2014), where the sketch size  $c$  is divisible by the number of blocks  $D^3$ . As shown in the last 4 columns in Table 1, sketching reduces the time complexity of reward imputation from  $O(MBd^2)$  to  $O(Mcd^2)$  for all  $M$  actions in one episode, where  $c < B$ . When  $Mc \approx B$ , the overall time complexity of our reward imputation using sketching is even comparable to that without reward imputation, i.e., a  $O(Bd^2)$  time complexity.
Updated Policy using Imputed Rewards. Inspired by the UCB strategy (Li et al., 2010), the updated policy for online decision of the  $(n + 1)$ -th episode can be formulated using the imputed rewards (parameterized by  $\bar{\theta}_A^{n + 1}$  in (2)) or the sketched version of imputed rewards (parameterized by  $\tilde{\theta}_A^{n + 1}$  in (6)). Specifically, for a new context  $s$ ,