Paper: Provable Benefits of Multi-task RL under Non-Markovian Decision Making Processes (ICLR-2024)

Reviewer comment (Reproducibility):
The generality of the theorems in the paper allows the various common structure with different $\eta$-bracketing number of the joint model class. Several examples are already explained in the paper such as the multi-task observable POMDP sharing the common transition kernel. The generic upstream algorithm is highly computational inefficient in building the confidence set and find the optimistic policy. Therefore, an important question is how the optimization steps in the algorithm look like in a specific setting (e.g., the multi-task POMDPs). This helps to evaluate the effectiveness of the upstream algorithms.

Author response:
**A2:** The optimization step in our algorithm (in Line 3) contains two components: (1) The optimization over the model parameter  $\boldsymbol{\theta},\boldsymbol{\theta}'$ subject to the constraints that define the confidence set $\boldsymbol{\mathcal{B}}_k$ in Line 7 of Algorithm 1. Note that the maximization in the constraints in Line 7 will need one inner-loop of optimization to solve to define the constraint thresholds. For a given specific setting, certain function approximations (e.g., neural networks) can be used to parameterize the model transitions in an efficient way. Then the standard constrained optimization algorithms such as primal-dual method can be applied to solve the problem. (2) The optimization over policy can be done by first parameterizing policy and then using policy gradient methods. Alternatively, some papers in practice use tree search and sample-based methods to solve policy optimization in POMDPs [5, 6, 7].

Reviewer comment (Theory):
It seems that the $\eta$-bracketing number is used to build the optimism of optimistic policies with uniform convergence over the bracketing set. This is essentially (and also used as) the covering number of the model class in terms of $\|\cdot\|^{\mathrm{p}}_{\infty}$ norm. Why use bracketing number as the name instead of calling it the covering number directly?

Author response:
**A3:** Great question! Bracketing number is more suitable in our study for providing two-sided bounds (upper and lower bounds) to control the total variation distance between model classes, but we don't see a direct way to use the covering number for serving this purpose. Specifically, the bracketing number is used for univariate- or vector-valued functions classes, which satisfies certain partial order relations. Based on the partial order relation, the bracketing set requires a "two-sided" and "bracketing type" covering to control both upper  and lower bounds of the elements from target class, which is useful in analysis to control the total variation distance. In contrast, the covering set does not provide  upper and lower bounds and only require being close to elements from target class. Mathematically, for a vector-valued functions $\mathbf{f}: \mathcal{X} \to \mathbb{R}^d$, the $\epsilon$-bracketing number requires a bracket $[\mathbf{l},\mathbf{g}]$ such that $\mathbf{l}\leq\mathbf{f}\leq\mathbf{g}$ and $\||\mathbf{l}-\mathbf{g}\||\leq \epsilon$. However, the $\epsilon$-covering number only requires a ball with center $\mathbf{r}$ such that $\||\mathbf{r}-\mathbf{f}\||\leq \epsilon$.

Reviewer comment (Presentation):
The biggest weakness of this paper is that it crams too much content in the main paper, and does not use the appendix to explicate it. This rushed discussion makes the job of reading the paper more difficult than it needs to be. For example, consider the assumptions of rank-$r$ and $\gamma$-well-conditioned PSRs. These two assumptions are present throughout the paper, but they never get their deserved attention. A mere half a page of terse definitions for them might be justified in the main paper due to page number constraints, but it is hard to justify why they were never given due discussion in the appendix. This discussion should discuss the intuitive meaning of these assumptions, examples of PSRs which satisfy the assumptions (otherwise questions like, is the set $\Theta$ even non-empty, surface), examples of PSRs which don't satisfy the assumptions, what fails in the proofs if each assumption is relaxed, etc.

Author response:
**A1:** Thanks for the careful reading. We humbly acknowledge that the content in the main text is rather dense. If this has hinder the reviewer's understanding of our work, we sincerely apologize. If this paper is accepted, we will streamline the paper to make sure it is more readable and provide intuitions for, e.g., the $\gamma$-well-conditioned PSRs, in the Appendix following the reviewer's suggestions. 

Briefly speaking, the   **$\gamma$-well condition** property of PSRs ensures that error amplification which results from estimating $\psi^*(\tau\_h)$ is not extremely large since otherwise,  there exists a hard instance that is not well-conditioned and learning a good policy requires an exponentially large sample complexity [1]. The definition of $\phi\_h$
 and $\mathbf{M}\_h$ can be found on page 4, between Eqns. (1) and (2). The definition of the norm $\\|\cdot\\|\_{\infty}^{\mathrm{p}}$ (which depends on policy $\mathrm{p}$) can be found in the 4th paragraph on page 5. We have ensured that all mathematical symbols are defined before use, although due to the density of results, they may not be very prominently displayed.

Reviewer comment (Writing):
The notation $\phi_h$ and $\mathbf{M} _ h$ is never even defined, and the norm $\lVert \cdot \rVert_\infty^p$ is "pulled out of a hat" without explanation.

Author response:
**A1:** Thanks for the careful reading. We humbly acknowledge that the content in the main text is rather dense. If this has hinder the reviewer's understanding of our work, we sincerely apologize. If this paper is accepted, we will streamline the paper to make sure it is more readable and provide intuitions for, e.g., the $\gamma$-well-conditioned PSRs, in the Appendix following the reviewer's suggestions. 

Briefly speaking, the   **$\gamma$-well condition** property of PSRs ensures that error amplification which results from estimating $\psi^*(\tau\_h)$ is not extremely large since otherwise,  there exists a hard instance that is not well-conditioned and learning a good policy requires an exponentially large sample complexity [1]. The definition of $\phi\_h$
 and $\mathbf{M}\_h$ can be found on page 4, between Eqns. (1) and (2). The definition of the norm $\\|\cdot\\|\_{\infty}^{\mathrm{p}}$ (which depends on policy $\mathrm{p}$) can be found in the 4th paragraph on page 5. We have ensured that all mathematical symbols are defined before use, although due to the density of results, they may not be very prominently displayed.

Reviewer comment (Experiments):
While not necessarily a weakness, it would have been nice to demonstrate the usefulness of the theory developed in the paper on some simple experiments. Additionally, it would be nice to calculate the computational complexity of implementing the algorithms. Note that these additions are not necessarily expected for this paper as it is already very terse.

Author response:
**A3:** Thanks for the kind suggestions. This paper is primarily positioned as a theoretical contribution. If we can find a suitable dataset, we will conduct  some experiments and report the results in the final version of the paper.

Reviewer comment (Novelty):
The downstream learning seems to be only applying previous results on a smaller downstream model class, without further new ideas.

Author response:
Thanks for the comments. In fact, our downstream analysis involves the following novel development that is not in previous results. Since the downstream model class is not known in advance, the agent needs to use information from the upstream to approximate the model, which introduces an approximation error $\epsilon_0$. Our handling of such an error here features a novel technique of using **Renyi divergence** to measure the approximation error, which guarantees the sub-optimality bound **without requiring the realizability condition**. Such a technique has not been developed in previous results.

Reviewer comment (Evaluation):
Can these results generalize beyond low-rank problems? (maybe not low-rank, but some other structures)

Author response:
Thanks for the question! A potential (more general) class of problems can be general decision making problems with small **generalized Eluder coefficient** (GEC) in [1]. We expect that it is possible to extend our multi-task framework to multi-task GEC type of problems. One key step is to come up with a new complexity measure that suitably particularizes to the Eluder dimension and captures the latent structure between tasks at the same time.

Reviewer comment (Theory):
The results hold for finite action and observation spaces. Can they be generalized to general infinite spaces? (maybe using function approximation or other techniques)

Author response:
We thank the reviewer for this question. Here are our initial thoughts. For generalization to infinite observation spaces, we can possibly leverage the techniques in [2], which  extended the finite observation setting to its continuous  counterpart. We expect that similar techniques would be applicable to this proposed extension. For generalization to infinite action spaces, 
a recent study under low-rank MDPs investigated continuous action spaces [3]. Similar to [3], with some smoothness assumptions, we expect that these techniques might be used  to generalize our results to infinite continuous action spaces.

Reviewer comment (Theory):
A set of tasks that is not observed uniformly. We still want uniformly good behavior on all the tasks, but the algorithm can only sample the tasks from a distribution, rather than go over each of them one by one in each iteration. This is highly relevant to how we train RL agents in POMDPs with procedural generation of the initial state, because procedural generation gives very coarse control over the resulting initial state distribution. Here, I am seeing each initial state in the support of the distribution as a task. The downstream learning is also interesting for this.

Author response:
**A1:** Thanks for the great questions and suggestions! Concerning the case that tasks are not observed uniformly, our current framework should be able handle this example with a slight change of the algorithm, namely, that  the confidence set should be modified so that it is  constructed using only the data observed. The data-collection step (Line 4) should be changed from ``collecting data for all task n`` to ``only collecting data for certain task n sampled from the prior distribution`` $\mathcal{D}_N$. Then, we only use the collected data to construct the confidence set (Line 7). In this way, we surmise that the learning goal should be modified to $\mathbb{E}\_{n}[V\_{\theta\_n^*, R\_n}^{\pi^*} - V\_{\theta\_n^*, R\_n}^{\pi\_n} ] \leq \epsilon$, where task $n$ follows the distribution related to $\mathcal{D}\_N$.

Reviewer comment (Theory):
A single PSR with block structure in the dynamics. This is like the example above, but the multiple tasks are not explicitly recognized as such.

Author response:
**A2:** Thanks for the example, which does not explicitly contain multiple tasks but the single PSR has additional structure that a suitably designed algorithm can possibly exploit. We envision that this setting will require more careful thinking and analyses because the algorithm first has to **identify** that the PSR indeed contains the additional structure. Even if this is known, the blocks must be uncovered for further exploitation. A possible solution is to design an additional mechanism which estimates the relationship between the current observation and its underlying environment. We envision that the nature of the result would be rather different. In particular, a quantity that is distinctly different from our $\eta$-bracketing number would emerge as a quantification of the utility of the block structure in the single PSR.

Reviewer comment (Experiments):
What impact does more and more training on the upstream tasks have on the zero-shot performance on the downstream task?

Author response:
**A3:** In the zero-shot case, since there is no interaction with the environment for the  downstream tasks, the performance will be  determined solely by the approximation error $\epsilon\_0$ of using upstream information to estimate shared structure with downstream model. To wit, the more the training data utilized by the upstream tasks, the smaller the $\epsilon\_0$, and the better the zero-shot performance. Since the upstream task may not fully capture all model parameters of the downstream task, in practice, if possible, few-shot interaction with the downstream environment is typically   recommended to improve the overall performance.

Reviewer comment (Theory):
What impact does more and more training on the upstream tasks have on the speed of learning on the downstream task? It would be a surprising and interesting find if some amount of "pre-training" upstream would actually improve the rate of convergence of the downstream. I guess it's more likely that the guarantee would be like "if you want to train for a given budget X downstream, then you can get good rates if you train for Y amount of experience upstream."

Author response:
**A4:** In some sense, our Theorem 2, which characterizes the benefits of downstream transfer learning provides such type of guarantee. Theorem 2 provides the total variation distance between estimated model and true model, which is a function of the log bracketing number $\sim \beta\_0$ and the approximation error from upstream learning $\epsilon\_0$. Specifically, the rate of learning (convergence speed) is 
inversely proportional to $\beta\_0$, which scales with $\epsilon\_0$. Similar to our answer to the above question, the more the training data of upstream tasks, the smaller the approximation error, and the faster the speed of learning.

Reviewer comment (Theory):
At a higher level, not as a criticism to the paper, though, I find the overall setting a bit odd. The proposed algorithm does not have any sequential interaction with the environment. Instead it runs the policy in the tasks for collecting data and updates its confidence set. What I find odd is therefore that the tasks can be so hard intra-episode that there is nothing we can do by adapting as we act, and we might as well pick policies, deploy them, and update. I guess the setting also does not quite apply to the kinds of environments that would be "easy" and where intra-episode adaptation could help improve performance.

Author response:
**A5:** Thanks for the comment! The reason that there is no intra-episode adaptation is that in the setting we consider which involves sequential decision making with a finite horizon, the transition kernel within an episode is **heterogeneous**. Hence, the data in the previous time steps with one episode does not provide any useful information in the subsequent time steps. As a result, there is no benefit  in  updating the policy for later steps using the earlier information. Instead, the policy update is made after each entire episode which does provide useful information for later episodes. We would like to further point out that in some ``easy`` case where the transitions share some similarities inside one episode, adapting the policy within the episode can improve performance.

Reviewer comment (Writing):
I am somewhat confused about what $\pi(\omega_h)$ means in Eq. 2, considering that $\pi$ for the end of the episode $\omega$ depends on what happened in the beginning of the episode ($\tau$). So how does $\tau$ factor into Eq. 2?

Author response:
**A6:** We apologize for the confusion and the error in Eq. (2) which should also be maximized over all possible $\tau_h$. Eq. (2) should be
$$
    \forall h \in [H], \max\_{x\in\mathbb{R}^{d\_h}:\|x\|_1\leq 1} \max\_{\pi\in\Pi}\max\_{\tau_h \in \mathcal{H}\_h} \sum\_{\omega\_h \in \Omega\_h}\pi(\omega\_h|\tau\_h)|\mathbf{m}(\omega\_h)^{\top}x|\leq \frac{1}{\gamma},
$$
where $\pi(\omega\_h|\tau\_h)=\frac{\pi(\omega\_h,\tau\_h)}{\pi(\tau\_h)}=\pi(a\_H|o\_H,a\_{H-1},\ldots,o\_{h+1},\tau\_h)\ldots\pi(a\_{h+1}|o\_{h+1},\tau\_h)$. We have clarified the definition in the revision version of the manuscript.

Reviewer comment (Theory):
It would also be nice to understand the relationship between r from Definition 1 and multiple tasks. Considering the correspondence between block-dynamics and multiple tasks I mentioned above (that is, multiple tasks can be put together into a single POMDP that samples the task as part of the initial state), what is the relationship between r and N? r is arguably harder to scrutinize than N and the shared structure between tasks, so maybe it's possible to get rid of r as a proxy for multiple tasks and formalize everything in multitask terms?

Author response:
**A7:** The rank $r$ measures, roughly speaking, the amount of correlation between different observations. A large $r$ indicates that there is less correlation between trajectories. However, $N$ is simply the number of potential tasks. If we combine multiple tasks into a single PSR, then $r$ measures how much information for a new task we can obtain from previously observed tasks. Therefore, $N$ is the number of tasks, while $r$ is a latent complexity measure that quantifies   the amount of correlations between tasks, similar to the $\eta$-bracketing number. The parameter $r$ should not be gotten rid of as it is a critical parameter of interest in our work as it quantifies the correlation among tasks.

Reviewer comment (Presentation):
I was a bit confused reading algorithm one because it refers to quantities that are defined after it is shown ($\nu^{\pi_n,k}$).

Author response:
**A8:** Thanks for the suggestion. We will adjust the algorithm so that it appears after the description. This font represents the whole environment. In terms of the "linear span", we agree that it would be better to use the notation of model dynamics $\mathbb{P}\_{\theta_1},\ldots,\mathbb{P}\_{\theta_m}$.

Reviewer comment (Presentation):
In Example 3, the P for the core tasks seem to be on the wrong font type?

Author response:
**A8:** Thanks for the suggestion. We will adjust the algorithm so that it appears after the description. This font represents the whole environment. In terms of the "linear span", we agree that it would be better to use the notation of model dynamics $\mathbb{P}\_{\theta_1},\ldots,\mathbb{P}\_{\theta_m}$.