SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Riccardo Poiani
DEIB, Politecnico di Milano riccardo.POiani@polimi.it
Nicole Nobili
DEIB, Politecnico di Milano  
nicole.nobili@mail.polimi.it
DEIB, Politecnico di Milano albertomaria.metelli@polimi.it
DEIB, Politecnico di Milano  
marcello/restelli@polimi.it
Policy evaluation via Monte Carlo (MC) simulation is at the core of many MC Reinforcement Learning (RL) algorithms (e.g., policy gradient methods). In this context, the designer of the learning system specifies an interaction budget that the agent usually spends by collecting trajectories of fixed length within a simulator. However, is this data collection strategy the best option? To answer this question, in this paper, we consider as quality index the variance of an unbiased policy return estimator that uses trajectories of different lengths, i.e., truncated. We first derive a closed-form expression of this variance that clearly shows the sub-optimality of the fixed-length trajectory schedule. Furthermore, it suggests that adaptive data collection strategies that spend the available budget sequentially might be able to allocate a larger portion of transitions in timesteps in which more accurate sampling is required to reduce the variance of the final estimate. Building on these findings, we present an adaptive algorithm called Robust and Iterative Data collection strategy Optimization (RIDO). The main intuition behind RIDO is to split the available interaction budget into mini-batches. At each round, the agent determines the most convenient schedule of trajectories that minimizes an empirical and robust estimate of the estimator's variance. After discussing the theoretical properties of our method, we conclude by assessing its performance across multiple domains. Our results show that RIDO can adapt its trajectory schedule toward timesteps where more sampling is required to increase the quality of the final estimation.
In Reinforcement Learning [RL, Sutton and Barto, 2018], an agent acts in an unknown, or partially known, environment to maximize/estimate the infinite expected discounted sum of an external reward signal, i.e., the expected return. Monte Carlo evaluation [MC, Owen, 2013] is at the core of many successful RL algorithms. Whenever a simulator with reset possibility is available to the learning systems designer, a large family of approaches [Williams, 1992, Baxter and Bartlett, 2001, Schulman et al., 2015, 2017, Cobbe et al., 2021] that can be used to solve the RL problem relies on MC simulations for estimating performance or gradient estimates on the task to be solved. In this scenario, since the goal is to estimate the expected infinite sum of rewards, the designer usually specifies a sufficiently large estimation horizon  $T$ , along with a transition budget  $\Lambda = QT$ , so that the agent
interacts with the simulator, via MC simulation, collecting a batch of  $Q$  episodes of length  $T$ . In this sense, the agent spends its available budget  $\Lambda$  uniformly along the estimation horizon.
In the context of MC policy evaluation, where the goal lies in estimating the performance of a given policy via MC simulations, Poiani et al. [2023] have recently shown that, given the discounted nature of the RL objective, this uniform-in-the-horizon budget allocation strategy may not be the best option. The core intuition behind their work is that, since rewards are exponentially discounted through time, early interactions weigh exponentially more than late ones, and, consequently, a larger portion of the available budget  $\Lambda$  should be dedicated to estimating the initial rewards. To theoretically validate this point, the authors designed a non-adaptive budget allocation strategy which, by exploiting the reset possibility of the simulator, leads to the collection of trajectories of different lengths, i.e., truncated. They show that this approach provably minimizes Höeffding-like confidence intervals [Boucheron et al., 2003] around the empirical estimates of the expected return. Remarkably, this implies a robustness w.r.t. the uniform strategy that holds for any pair of environment and policy to be evaluated, thus, clearly establishing the theoretical benefits of the proposed method.
Nevertheless, it has to be noticed that although minimizing confidence intervals around the expected return estimator comes with desirable theoretical guarantees (e.g., PAC-bound improvements [Even-Dar et al., 2002]), the resulting schedule of trajectories is computed before the interaction with the environment (being determined by the discount factor). Consequently, as the usual uniform-in-the-horizon scheme, it fails to adapt to the peculiarities of the problem at hand, and, ultimately, might not produce a low error estimate. For the sake of clarity, we illustrate this sub-optimality of pre-determined schedule of trajectories with the following extreme examples.
Example 1. Consider an environment where the reward is gathered only at the end of the horizon  $T$  (e.g., a goal-based). In this scenario, any strategy that truncates trajectories is intuitively sub-optimal, and we expect that an intelligent agent will spend all its budget according to the uniform schedule.
Example 2. Conversely, consider a problem where the reward is different from 0 in the first interaction step only (e.g., in the case of a highly sub-optimal policy that immediately reaches the "zero reward region" of an environment); the uniform schedule wastes a significant portion of its budget collecting samples without variability, and, to reduce the estimation error, we would like the agent to spend all of its interaction budget estimating the reward of the first action.
Abstracting away from the previous examples, we realize that the main issue of existing approaches arises from the fact that determining a schedule of trajectories before interacting with the environment does not allow the agent to adapt it to the environment peculiarities, allocating more samples where this is required to obtain a high-quality estimate. For this reason, in this work, we focus on designing adaptive data collection strategies that aim directly at minimizing the error of the final estimate. Our main intuition lies in splitting the available budget  $\Lambda$  into mini-batches and adapting online the data collection strategy of the agent based on the previously collected information.
Original Contributions and Outline After introducing the necessary notation and backgrounds (Section 2), we consider the problem of maximizing the estimation quality of a policy expected return estimator using trajectories of different lengths collected via MC simulation with a finite budget  $\Lambda$  of transitions (Section 3). More specifically, since we use an unbiased return estimator, we consider its variance as a quality index, of which we derive a closed-form expression and analyze it for every possible schedule of trajectories. Then, we define the optimal trajectories schedule as the one that attains the minimum variance subject to the available budget constraint. As expected, computing this optimal data collection strategy requires knowledge of the underlying environment (e.g., the variance of the rewards at each timestep), which is not available to the agent prior to the interaction. Nevertheless, as we shall see, all the quantities that define the optimal strategy can be estimated from the data. These facts confirm our intuition about the weakness of non-adaptive schedules of trajectories and suggest that algorithms that spend the available budget  $\Lambda$  iteratively might be able to dynamically allocate their budget to minimize the variance of the final estimate. Building on these findings, in Section 4, we present a novel algorithm, Robust and Iterative Data collection strategy Optimization (RIDO), which splits its available budget  $\Lambda$  into mini-batches of interactions that are allocated sequentially to minimize an empirical and robust estimate of the
<sup>1</sup>While another large class of RL algorithms is based on Temporal Difference [TD, Sutton and Barto, 2018] learning, which do not require the finite horizon nor the reset possibility, Monte Carlo simulation approaches continue to be extensively adopted. Indeed, unlike TD methods, they can be applied effortlessly to non-Markovian environments, which is a common occurrence in real-world problems.
objective function of interest, i.e., the variance of the estimator. Furthermore, we perform a statistical analysis on the behavior of RIDO, and we derive theoretical guarantees expressed as upper bounds on the variance of the policy return estimator. Our result shows that, under favorable conditions, the variance of the return estimator computed by RIDO is of the same order as the one of the oracle's baseline. To conclude, in Section 5, we conduct an experimental comparison between RIDO and non-adaptive schedules. As we verify, our method achieves the most competitive performance across different domains, discount factor values, and budget, thus clearly highlighting the benefits of adaptive strategies over pre-determined ones.
2 Backgrounds and Notation
This section provides the notation and necessary backgrounds used in the rest of this document.
Markov Decision Processes A discrete-time Markov Decision Process [MDP, Puterman, 2014] is defined as a tuple  $\mathcal{M} \coloneqq (\mathcal{S}, \mathcal{A}, R, P, \gamma, \nu)$ , where  $\mathcal{S}$  is the set of states,  $\mathcal{A}$  is the set of actions,  $R: \mathcal{S} \times \mathcal{A} \to [0,1]$  is the reward function the specifies the reward  $R(s,a)$  received by the agent upon taking action  $a$  in state  $s$ ,  $P: \mathcal{S} \times \mathcal{A} \to \Delta(\mathcal{S})^2$  is the transition kernel that specifies the probability distribution over the next states  $P(\cdot | s, a)$ , when taking action  $a$  in state  $s$ ,  $\gamma \in (0,1)$  is the discount factor, and  $\nu \in \Delta(\mathcal{S})$  is the distribution over initial states. The agent's behavior is modeled by a policy  $\pi: \mathcal{S} \to \Delta(\mathcal{A})$ , which for each state  $s$ , prescribes a distribution over actions  $\pi(\cdot | s)$ . A trajectory  $\tau_h$  of length  $h$  is a sequence of states and actions  $(s_0, a_0, s_1, \ldots, s_{h-1}, a_{h-1}, s_h)$  observed by following  $\pi$  for  $h$  steps, where  $s_0 \sim \nu$ ,  $a_t \sim \pi(\cdot | s_t)$ , and  $s_{t+1} \sim P(\cdot | s_t, a_t)$  for  $t < h$ . The return of a trajectory is defined as  $G(\tau_h) = \sum_{t=0}^{h-1} \gamma^t R_t$ , where  $R_t$  is shortcut for  $R(s_t, a_t)$ . The agent that is following policy  $\pi$  is evaluated according to expected cumulative discounted sum of rewards over an estimation horizon  $T$ , namely  $J(\pi) = \mathbb{E}_{\pi} \left[ \sum_{t=0}^{T-1} \gamma^t R_t \right]$ , where the expectation is taken w.r.t. the stochasticity of the policy, the transition kernel, and the initial state distribution.
Data Collection Strategy Poiani et al. [2023] formalized the concept of Data Collection Strategy (DCS) to model how the agent collects data within an environment. More specifically, given an interaction budget  $\Lambda \in \mathbb{N}$  such that  $\Lambda \mod T = 0$ , a DCS is defined as a  $T$ -dimensional vector  $\boldsymbol{m} \coloneqq (m_1, \dots, m_T)$  where  $m_h \in \mathbb{N}$  and  $\sum_{h=1}^{T} m_h h = \Lambda$ . Each element  $m_h$  specifies the number of trajectories of length  $h$  that the agent collects in the environment while following a policy  $\pi$ . Given a DCS  $\boldsymbol{m}$ , it is possible to compute the total number of steps  $\boldsymbol{n} \coloneqq (n_0, \dots, n_{T-1})$  that will be gathered by the agent at any step  $t$ ; more specifically, the following relationship holds:  $n_{T-1} = m_T$ , and  $n_t = n_{t+1} + m_{t+1}$  for  $t < T-1$ . For this reason, in the rest of the paper we will adopt the most convenient symbol depending on the context. For any DCS  $\boldsymbol{m}$  such that  $m_T \geq 1$  holds, it is possible to build the following unbiased estimator of  $J(\pi)$ :
$$
\hat {J} _ {\boldsymbol {m}} (\pi) = \sum_ {h = 1} ^ {T} \sum_ {i = 1} ^ {m _ {h}} \sum_ {t = 0} ^ {h - 1} \frac {\gamma^ {t}}{n _ {t}} R _ {t} ^ {(i)}. \tag {1}
$$
The two external summations in Equation (1) sum over the collected trajectories of different lengths a rescaled empirical trajectory return, where the reward at step  $t$  is divided by the number of samples collected at step  $t$ . In the case the budget  $\Lambda$  is spent uniformly, i.e.,  $\boldsymbol{m} = \left(0, \dots, 0, \frac{\Lambda}{T}\right)$ , Equation (1) reduces to the usual Monte Carlo estimator of  $J(\pi)$ , namely  $\frac{T}{\Lambda} \sum_{i=1}^{\Lambda/T} \sum_{t=0}^{T-1} \gamma^t R_t^{(i)}$ .
Robust Data Collection Strategy Optimization Leveraging the estimator of Equation (2), Poiani et al. [2023] investigated alternatives to the usual uniform-in-the-horizon DCS from the worst-case perspective of confidence intervals [Boucheron et al., 2003]. More specifically, given  $m$  such that  $m_T \geq 1$ , the estimator of Equation (1) enjoys the following generalization of the Hoeffding
2Given a set  $\mathcal{X}$  , we denote with  $\Delta (\mathcal{X})$  the set of probability distributions over  $\mathcal{X}$
<sup>3</sup>As common in Monte-Carlo simulation [see e.g., Papini et al., 2022] we approximate the infinite horizon MDP model with a finite estimation horizon  $T$ . Indeed, if  $T$  is sufficiently large, i.e.,  $T = \mathcal{O}\left(\frac{1}{1 - \gamma} \log \frac{1}{\epsilon}\right)$ , the expected return computed with horizon  $T$  is  $\epsilon$  close to the infinite-horizon one [Kakade, 2003].
4Rescaling by  $n_t$  prevents the estimate from being biased towards time steps for which more samples are available.
confidence intervals holding with probability at least  $1 - \delta$ :
$$
| J (\pi) - \hat {J} _ {\boldsymbol {m}} (\pi) | \leq \sqrt {\frac {1}{2} \log \left(\frac {2}{\delta}\right) \sum_ {t = 0} ^ {T - 1} \frac {d _ {t}}{n _ {t}}}, \tag {2}
$$
where  $d_{t} = \frac{\gamma^{t}\left(\gamma^{t} + \gamma^{t + 1} - 2\gamma^{T}\right)}{1 - \gamma}$  controls the relative importance of samples gathered at step  $t$ . Poiani et al. [2023] designed a closed-form DCS that provably minimizes the bound of Equation (2). Since  $d_{t}$  is a decreasing function of time whose decay speed is governed by the discount factor  $\gamma$ , the aforementioned DCS gives priority to the collection of experience at earlier time steps, i.e., it truncates the trajectories. Note that the smaller  $\gamma$ , the higher the number of samples reserved for earlier time steps. We refer the reader for Theorem 3.3 and Theorem B.10 of their work for the exact expressions of the resulting robust DCS. However, we remark that the resulting schedule is non-adaptive (i.e., it is computed before the interaction with the environment takes place) and its shape depends exclusively on  $\Lambda$ ,  $\gamma$ , and  $T$ .
3 Toward Adaptive Data Collection Strategies
In this section, we lay down the theoretical groundings behind optimizing data collection strategies that directly aim at minimizing the final estimation error. We stick to methods that adopt the estimator of Equation (1), which has a simple interpretation and desirable theoretical properties. More specifically, since the estimator is unbiased, the Mean Squared Error (MSE) simply reduces to the variance. For this reason, to set a proper baseline for DCS optimization (i.e., an optimal strategy according to the MSE), we first analyze the variance of the estimator of Equation (1) for an arbitrary DCS  $m$ . The following Theorem (proof in Appendix A) summarizes our result.
Theorem 3.1. Consider a generic  $DCS\mathbf{m}$  such that  $m_T \geq 1$ , then:
$$
\mathbb {V} \operatorname {a r} _ {\boldsymbol {m}} \left[ \hat {J} _ {\boldsymbol {m}} (\pi) \right] = \sum_ {t = 0} ^ {T - 1} \frac {1}{n _ {t}} \left(\gamma^ {2 t} \mathbb {V} \operatorname {a r} \left(R _ {t}\right) + 2 \sum_ {t ^ {\prime} = t + 1} ^ {T - 1} \gamma^ {t + t ^ {\prime}} \mathbb {C} \operatorname {o v} \left(R _ {t}, R _ {t ^ {\prime}}\right)\right) =: \sum_ {t = 0} ^ {T - 1} \frac {f _ {t}}{n _ {t}}. \tag {3}
$$
Theorem 3.1 expresses, in closed form, the variance of the estimator of Equation (1) when adopted with an arbitrary DCS that guarantees the estimation to be unbiased (i.e.,  $m_T \geq 1$ ). From Equation (3), we can see that this variance results in a summation, over the different time steps, of  $\frac{1}{n_t}$ , i.e., the reciprocal of the number of samples collected under  $m$  at step  $t$ , multiplied by the variance of the reward at step  $t$  plus the covariances between  $R_t$  and the rewards gathered at future steps. For brevity, we shortcut this term with  $f_t$ . Furthermore, Theorem 3.1 leads to a direct formulation of an optimal DCS baseline for our setting. More specifically, given a budget  $\Lambda$ , we define the optimal DCS  $n^*$  for the estimator in Equation (1) as the solution of the following optimization problem:
$$
\min  _ {\boldsymbol {n}} \quad \sum_ {t = 0} ^ {T - 1} \frac {1}{n _ {t}} \left(\gamma^ {2 t} \mathbb {V} \operatorname {a r} \left(R _ {t}\right) + 2 \sum_ {t ^ {\prime} = t + 1} ^ {T - 1} \gamma^ {t + t ^ {\prime}} \mathbb {C} \operatorname {o v} \left(R _ {t}, R _ {t ^ {\prime}}\right)\right)
$$
$$
\text {s . t .} \quad \sum_ {t = 0} ^ {T - 1} n _ {t} \leq \Lambda \tag {4}
$$
$$
n _ {t} \geq n _ {t + 1}, \quad \forall t \in \{0, \dots , T - 2 \}
$$
$$
n _ {t} \in \mathbb {N} _ {+}, \quad \forall t \in \{0, \dots , T - 1 \},
$$
where the constraints  $n_t \geq n_{t+1}$  directly encode the sequential nature of the interaction with the environment. At this point, some comments are in order. First of all, we notice that the above optimization problem nicely captures the intuitive examples of Section 1.
Example 1 (cont.). When the reward is different from 0 in the last interaction step only, the objective function reduces to  $\frac{\gamma^{2(T-1)}}{n_{T-1}}\mathbb{Var}\left[R_{T-1}\right]$ , which is clearly minimized for the uniform strategy.
Example 2 (cont.). Conversely, when the reward is different from 0 in the first step only, we obtain  $\frac{\operatorname{Var}[R_0]}{n_0}$ , meaning that the entire interaction budget should be dedicated to estimate  $R_0$ .
Algorithm 1 Robust and Iterative DCS Optimization (RIDO).
Require: Interaction budget  $\Lambda$ , batch size  $b$ , robustness level  $\beta$ , policy  $\pi$   
1: Collect  $\mathcal{D}$  using policy  $\pi$  and  $\hat{n}_0 = \left(\frac{b}{T}, \dots, \frac{b}{T}\right)$   
2: Set  $K = \frac{\Lambda}{b}$  and initialize empirical estimates  $\sqrt{\widehat{\mathrm{Var}}_1[R_t]}$  and  $\widehat{\mathrm{Cov}}_1[R_t, R_t']$   
3: for  $i = 1, \dots, K - 1$  do  
4: Collect  $\mathcal{D}_i$  using policy  $\pi$  and  $\hat{n}_i$ , where  $\hat{n}_i$  is computed solving problem (5)  
5: Update empirical estimates  $\sqrt{\widehat{\mathrm{Var}}_i[R_t]}$  and  $\widehat{\mathrm{Cov}}_i[R_t, R_t']$  using  $\mathcal{D}_i$  and set  $\mathcal{D} \gets \mathcal{D} \cup \mathcal{D}_i$   
6: end for
Furthermore, the formulation of problem (4) highlights the pitfalls of a non-adaptive data collection strategies. Indeed, consider, for the sake of clarity, the two examples mentioned above. Before executing policy  $\pi$ , the agent has no way of distinguishing between the two different objective functions, i.e.,  $\frac{\gamma^{2(T-1)}}{n_{T-1}} \mathbb{Var}[R_{T-1}]$  and  $\frac{\mathbb{Var}[R_0]}{n_0}$ , and, consequently, any pre-determined schedule fails to adapt to the actual objective function. More generally, this is because the optimal strategy resulting from the optimization problem (4) can be computed prior to the interaction with the environment only by an oracle that knows in advance the underlying reward process induced by the agent's policy  $\pi$  in the MDP. Nevertheless, we note that all the terms appearing in the objective function, i.e., the only unknowns in optimization problem (4), can be estimated if some interactions with the environment are available to the agent. This suggests that strategies that sequentially allocate the available budget  $\Lambda$  might successfully adapt their DCS to minimize Equation (3), i.e., the variance of the return estimator.
4 Robust and Iterative DCS Optimization
Given the findings of Section 3, we now present our algorithmic solution that aims at avoiding the highlighted pitfalls of pre-determined DCSs. Our approach is called Robust and Iterative Data collection strategy Optimization (RIDO), and its pseudocode is available in Algorithm 1. The central intuition behind RIDO lies in splitting the available budget  $\Lambda$  into mini-batches of interactions that the agent will allocate sequentially. At each iteration, the agent will compute the most convenient schedule of trajectories that optimizes an empirical and robust version of the objective function presented in (4), whose quality improves as the agent gathers more data.
We now describe in-depth the behavior of the algorithm. For simplicity of exposition and analysis, we suppose that the size of the mini-batch  $b$  is such that  $b \mod T = 0$  and  $b \geq 2T$ . At the beginning (Lines 1-2 in Algorithm 1), the agent spends the first mini-batch  $\hat{n}_0$  at collecting  $\frac{b}{T}$  trajectories of length  $T$  (i.e., the uniform approach). This preliminary collection phase is a starting round in which some initial experience is gathered to properly initialize estimates of relevant quantities used throughout the algorithm. More specifically, at each iteration  $i$ , the agent maintains empirical estimates of the unknown quantities that define the variance of the estimate, i.e., the standard deviation of the reward at step  $t$ , namely  $\sqrt{\widehat{\mathbb{V}}\mathrm{ar}_i[R_t]}$ , and the covariances between rewards at different steps, namely  $\widehat{\mathbb{C}}\mathrm{ov}_i(R_t,R_{t'})$ . Then, at each round (Lines 4-5 in Algorithm 1), the DCS of the current mini-batch  $\hat{n}_i$  is computed solving the optimization problem (5) whose objective function is a robust estimate of the objective function of the original optimization problem (4). More specifically, at each round  $i$ , the agent aims at solving the following problem:
$$
\min  _ {\boldsymbol {n}} \quad \sum_ {t = 0} ^ {T - 1} \frac {1}{n _ {t}} \left(\gamma^ {2 t} \left(\sqrt {\hat {\mathbb {V}}} \operatorname {a r} _ {i} (R _ {t}) + \mathbf {C} _ {i, t} ^ {\sigma}\right) ^ {2} + 2 \sum_ {t ^ {\prime} = t + 1} ^ {T - 1} \gamma^ {t + t ^ {\prime}} \left(\widehat {\mathbb {C}} \operatorname {c o v} _ {i} (R _ {t}, R _ {t ^ {\prime}}) + \mathbf {C} _ {i, t, t ^ {\prime}} ^ {c}\right)\right)
$$
$$
\text {s . t .} \quad \sum_ {t = 0} ^ {T - 1} n _ {t} \leq b \tag {5}
$$
where  $\mathbf{C}_{i,t}^{\sigma}$  and  $\mathbf{C}_{i,t,t'}^c$  are exploration bonuses for variances and covariances respectively, defined as:
$$
\mathrm {C} _ {i, t} ^ {\sigma} := \sqrt {\frac {2 \log (\beta)}{\sum_ {j = 1} ^ {i - 1} \hat {n} _ {j , t}}}, \quad \mathrm {C} _ {i, t, t ^ {\prime}} ^ {c} := 3 \sqrt {\frac {2 \log (\beta)}{\sum_ {j = 1} ^ {i - 1} \hat {n} _ {j , t ^ {\prime}}}}, \tag {6}
$$
where  $\beta \geq 1$  is a hyper-parameter that specifies the amount of exploration used to solve the optimization problem, and  $\hat{n}_{j,t}$  is the number of samples collected by RIDO during phase  $j$  at time step  $t$ . We now provide further explanations on the optimization problem (5) and Equation (6). First of all, we notice how each term in the original objective function, namely  $f_{t}$ , is replaced with its relative empirical estimation plus exploration bonuses, each of which is directly related to components