Transfer Q-Learning with Composite MDP Structures

Thank you for your insightful comments!

Low-Rank: We agree that clarifying this terminology is important for improving readability and avoiding confusion. In classical low-rank MDPs literature, "low-rank" refers to embedding-based models where the effective feature dimension $d$ is small compared to $|\mathcal{S}|$. In contrast, in our setting, we work in a high-dimensional regime with $p, q \gg N$, so $M^*$ is not globally low-rank. Prior methods developed for low-rank MDPs—which assume a small feature dimension—do not apply directly, since they cannot consistently estimate $M^\ast$ within $N$ trajectories.

To resolve this issue, we assume $M^* = L^* + S^*$, where $L^*$ is of low-rank $r$ and $S^*$ is sparse. This structure enables consistent estimation in high dimensions and is key to our theoretical analysis. Crucially, this assumption enables us to design an estimator that achieves minimax-optimal error rates for recovering $L^\ast$, $S^\ast$, and consequently $M^\ast$, independently of the ambient dimension $d = \max (p, q)$.

We have revised the manuscript to clarify the distinction between embedding-based low-rank MDPs and our structured high-dimensional setting.

Prior Work [1,2]: Our model can be written in the form of [1] as $P(s' \mid s, a) = \phi(s,a)^\top \alpha(s')$ with $\alpha(s') := (L^* + S^*)\psi(s')$. This reduces to [1] when $S^* = 0$. Similarly, our model matches the linear factored MDP in [2] via a Kronecker formulation: $P(s' \mid s,a) = (\phi(s,a) \otimes \psi(s'))^\top \operatorname{vec}(L^* + S^*)$. While structurally related, a key difference lies in estimation. Our setting allows $p, q \gg N$, and our estimator achieves minimax-optimal error rates that do not scale with $d = \max(p, q)$. In contrast, [1,2] assume low-dimensional or identifiable parameter spaces and are not suitable for high-dimensional regimes. Our structured assumption is especially crucial in transfer RL, enabling effective knowledge sharing via $L^*$ while capturing task-specific deviations via $S^*$.

Comp. Efficiency of $K_\psi$: We agree this is important. When $|\mathcal{S}|$ is large or infinite, we compute $K_\psi$ using a Monte Carlo approximation: $\hat{K}\_\psi = \frac{1}{m} \sum\_{i=1}^m \psi(s\_i)\psi(s\_i)^\top, \quad s\_i \sim \text{Unif}(\mathcal{S}).$ This is a standard approach in randomized numerical linear algebra [3] and is computed only once before online learning. Our method is thus comparable in efficiency to [1], which stores empirical covariances, but we use $K_\psi$ to build confidence regions specific to our model.

Clarification on Theorem 3.6: Our method handles the high-dimensional setting by explicitly exploiting the low-rank-plus-sparse structure of the transition core matrix $M^* = L^* + S^*$, where $L^*$ is of low-rank $r << d$ and $S^*$ is of sparsity $s << d$. This structure enables consistent estimation in regimes where $p, q \gg N$, and is key to our theoretical guarantees. We have updated Remark 3.7 to clearly state the point.

Crucially, our estimation procedure achieves minimax-optimal error rates for $L^*$, $S^*$, and $M^*$ that are independent of the ambient dimension $d = \max(p, q)$. We do not assume that the rank $r$ or sparsity level $s$ scale with $d$; under Assumption 3.1, both can remain small (e.g., constant or logarithmic in $d$), ensuring the estimation error bound does not grow with $d$ even in high-dimensional settings and our transfer learning result is explicitly free of $d$, demonstrating that the benefit of structural transfer carries over without dependence on the ambient dimension.

Technical Note and Transfer RL: Our regret analysis relies on bounding Frobenius norm errors and carefully transferring them to Bellman recursion. Improvements using variance reduction or stronger structural assumptions are possible future work. A major contribution of our paper is in transfer RL: our low-rank-plus-sparse structure enables sample-efficient transfer, capturing shared structure via $L^*$ and task-specific deviations via $S^*$. Prior work like FLAMBE [Agarwal et al., 2020] does not handle this generality. We have added clarifications throughout the revised paper.

We thank the reviewer again for your support and helpful comments, and we hope the improvements further affirm the strength and clarity of our submission.

References:

[1] Sample-Optimal Parametric Q-Learning Using Linearly Additive Features
[2] Model-Based Reinforcement Learning with Value-Targeted Regression
[3] Drineas & Mahoney, RandNLA: randomized numerical linear algebra, CACM 2016

Response to Reviewer fH8W

We sincerely thank the reviewer for the thoughtful feedback and recognition of our theoretical contributions. We address the two concerns raised below.

Q1: Lack of empirical validation of UCB-TQL

We appreciate the reviewer’s point. While this submission does not include empirical evaluations, our goal is to provide a theoretical foundation for transfer reinforcement learning (RL) in high-dimensional settings with structured heterogeneity. Specifically, UCB-TQL is the first algorithm with provable regret guarantees under the composite MDP model that includes both low-rank and sparse deviations, and achieves $\widetilde{\mathcal{O}}(\sqrt{e} H^5 N)$ regret that is independent of ambient dimension.

This setting is motivated by real-world applications (e.g., autonomous driving, robotics, and healthcare), where tasks often share underlying dynamics but differ in subtle, structured ways. In such domains, the transition dynamics may be high-dimensional, and assuming pure low-rank or dense similarity across tasks can be too simplistic. The low-rank-plus-sparse decomposition we propose reflects this nuanced task structure and provides an interpretable and tractable way to model variation.

Although we do not present experiments here, we see our work as a theoretical counterpart to recent empirical transfer RL methods. Our regret guarantees offer insights into when and how transfer should help, which is often unclear in empirical work. We believe our analysis will inspire future algorithms that combine empirical effectiveness with theoretical robustness, and we are currently implementing UCB-TQL in benchmark settings.

We hope the community will view our contribution as laying the groundwork for provably efficient transfer in high-dimensional RL, and we are excited to follow up with empirical validation in future work.

Q2: Limitation of sparsity-based task differences

We fully agree that real-world tasks can exhibit more complex deviations than those modeled by sparse differences. Our current work makes a structured but generalizable starting assumption: tasks differ sparsely in transition dynamics after aligning shared core structure. This assumption enables us to design statistically and computationally efficient estimators, while still covering rich task classes.

We appreciate the suggestion and have added discussion in the revised manuscript exploring possible extensions:

• Task-specific low-rank components: One natural direction is to allow each task to have its own low-rank variation on top of a shared component. This introduces additional complexity but could capture broader relationships.
• Decomposable low-rank spaces: Another extension is to learn a union of low-rank subspaces across tasks, leveraging techniques from subspace clustering and matrix factorization.
• Feature selection on $\Phi$ and $\Psi$: When feature sets are known but high-dimensional, incorporating feature selection or structured sparsity could adaptively identify task-relevant components.
• Group-sparse or structured differences: Beyond elementwise sparsity, our model could be extended to handle block or group-structured variations, which would better capture correlated shifts.

We view our work as a first step toward unifying high-dimensional structure with transfer learning, and we are enthusiastic about these extensions.

Closing Remarks

We again thank the reviewer for the constructive comments. We believe our theoretical framework makes a timely and valuable contribution to the transfer RL literature, particularly in advancing the understanding of when and how transfer helps in high-dimensional environments. By addressing the above concerns, we hope to improve the clarity, flexibility, and practical relevance of our work. We respectfully hope the reviewer may consider increasing their score.

We thank the reviewer for their positive assessment and constructive suggestions. We are especially grateful for the recognition of our theoretical contribution, which we believe offers a timely and foundational advancement in transfer reinforcement learning.

Our work provides the first provable regret guarantees for transfer RL in high-dimensional settings where the transition dynamics exhibit structured heterogeneity. By modeling transitions as $M^* = L^* + S^*$—with $L^*$ low-rank and $S^*$ sparse—we capture realistic task variations (e.g., shared core dynamics with sparse task-specific shifts). This structure allows us to derive minimax-optimal error rates for estimating $L^*$ and $S^*$, and achieve regret bounds that are independent of the ambient dimension $d = \max(p, q)$. These insights lay a theoretical foundation that complements ongoing empirical advances and help answer the fundamental question: when and how does transfer learning help in high-dimensional RL?

We have revised the manuscript to improve clarity in notation and address specific concerns raised:

• Section 2 – Notation of $\mathcal{P}$ and $\mathcal{P}V^\pi_{h+1}(s,a)$:
  $\mathcal{P}$ is a linear operator mapping a function $V : \mathcal{S} \to \mathbb{R}$ to a function over $(s,a)$ via $\mathcal{P}V^\pi\_{h+1}(s,a) = \sum\_{s'} \mathcal{P}(s'|s,a) V^\pi\_{h+1}(s')$. We have clarified this in the updated text.

• Line 154 – $[H]$ Notation:
  Yes, you are right. $[H]$ denotes the index set $\{1,2,\dots,H\}$. We have clarified this in the updated text.

• Line 191 – Summation in $K_\psi$ and Equation (1):
  The equality in Equation (1) is valid. $K_\psi$ is defined as the sum over all $s' \in \mathcal{S}$ and is independent from the outer summation over $s' \in \mathcal{S}$ in Equation (1). The simplification extracts common factors based on our composite MDP formulation.

• Line 210 – $\|\cdot\|_{2,\infty}$ Notation:
  This denotes the 2-to-infinity operator norm, i.e., the maximum $\ell_2$ norm of the rows. We've added this to the notation section for clarity.

• Equation (3) – Frobenius Norm:
  $\|\cdot\|_F$ refers to the Frobenius norm and was defined in the supplementary material, now clarified in the main text.

• Assumption 3.1 – Meaning of $\mu$ and $r$:
  Thank you for your careful reading! $\mu$ is the incoherence parameter (a constant > 1), and $r$ is the rank of the low-rank component $L^*$, now clarified in the main text.

Regarding examples of our model, consider personalized healthcare: the low-rank component captures typical patient responses to treatments, while the sparse component models anomalies or individual-specific deviations. Other real-world domains include recommendation systems and robotics, where core dynamics are shared but task-specific noise or variation exists. A counterexample would be settings where task variation arises from structured transformations or complex nonlinear mappings—scenarios that go beyond additive sparse deviations and are natural directions for future work.

We thank the reviewer again for their support and helpful comments, and we hope the improvements further affirm the strength and clarity of our submission.