Ask to Learn: Query-Based Collaborative Reinforcement Learning

Introduction

Collecting new trajectories for reinforcement learning (RL) can be costly, slow, or unsafe. In healthcare, for example, it is generally infeasible to deploy exploratory treatment policies on patients solely for data collection. Offline RL avoids such interaction by learning from historical records, but its guarantees depend critically on the coverage of the available data [17, 19, 31, 34]. The absence of feasible online exploration, however, does not preclude access to additional information through simulation. We study how a learner can acquire this information efficiently through collaboration.

We introduce query-based collaborative reinforcement learning, in which a learner augments a fixed offline dataset with observations from a collaborator-mediated simulator. At each query, the learner selects a state-time pair; the collaborator samples an action according to a fixed, possibly suboptimal policy, and the simulator returns the corresponding reward and next state. The learner thus controls where information is collected, but not which action generates it. Our central question is:

Given the available offline data, when is it possible to identify a near-optimal policy through collaboration, and how should queries be allocated to do so efficiently?

This setting is motivated by systems that couple decision models with simulators. In head-and-neck cancer treatment, Wentzel et al. [ [30] combine a model of physician treatment decisions with a patient simulator that predicts treatment responses and clinical outcomes. In our abstraction of such a system, a query specifies the patient state and treatment stage, while the physician model selects the treatment and the simulator supplies its consequences. Related structures arise in power-system dispatch, where a rule-based decision strategy interacts with a simulated grid [32], and in process control, where a human digital twin emulates an operator’s corrective actions in a simulated plant [6].

The collaborator need not be uniformly informative. For example, a hospital’s records may contain treatments that an externally developed clinical digital twin rarely or never selects, while the twin may provide information about patient states poorly represented in those records. The two sources should therefore be viewed as complementary rather than interchangeable. Unlike imitation learning, our objective is not to recover the collaborator’s policy, but to use both information sources to identify a near-optimal policy for the underlying decision process.

This division of control also distinguishes our setting from standard trajectory-based hybrid RL, which supplements offline data by executing exploration policies [7, 15, 26]. Here, the learner can request observations at a chosen state-time pair without first reaching that state through a trajectory. Conversely, unlike unrestricted generative-model access [1, 5, 25, 35], it cannot directly select the sampled action. The query strategy must therefore account for both the information already available offline and the collaborator’s sampling restrictions.

Uniformly querying all state-time pairs is generally inefficient. The offline data may already estimate some parts of the decision process accurately, while leaving others uncertain. More importantly, not every uncertainty is equally relevant: identifying a good policy requires resolving comparisons between plausible alternatives, rather than estimating every reward and transition equally accurately. Poor coverage alone therefore need not make a state the best query location. An efficient strategy should direct queries toward the policy comparisons that remain statistically difficult after accounting for the offline data, subject to which observations the collaborator can supply.

The sampling restriction also means that some uncertainty may persist indefinitely. If the collaborator never selects an action at a given state-time pair, no amount of querying can provide additional observations of that action there. Its estimation remains limited by the fixed offline data, potentially imposing an irreducible certification limit. A learning procedure must therefore account not only for how to allocate queries, but also for when additional queries cannot sufficiently reduce the uncertainty needed to certify the target accuracy.

Contributions. We initiate the framework of query-based collaborative reinforcement learning, where an instance-dependent analysis is natural because the value of additional queries depends strongly on the offline coverage, collaborator support and the structure of the underlying MDP, which can vary substantially across instances. Hence, a worst-case analysis obscures this heterogeneity and cannot capture how the offline data and collaborator jointly reduce the remaining uncertainty. For a target suboptimality $\varepsilon> 0$ and failure probability $\delta\in(0,1)$, our main results are as follows.

• Information-theoretic lower bound and local characterization. We derive an exact KL-based lower-bound benchmark for the additional collaborator queries required to identify an $\varepsilon$-optimal policy with probability at least $1-\delta$. We then derive a tractable local quadratic surrogate that makes explicit the roles of policy-comparison margins, local reward-transition variability, offline coverage, and collaborator support.

• Lower-bound-guided adaptive querying. We show that the same local variance structure governs the lower-bound approximation and the leading finite-sample uncertainty of policy-value comparisons. Guided by this connection, we develop an epoch-based algorithm that progressively narrows the candidate policy class and allocates queries through a local minimax design problem. The design jointly selects a reference policy and state–time query locations, concentrating observations on comparisons that remain difficult after exploiting the offline data.

• Finite-sample guarantees and near-matching complexity. We establish finite-query termination and high-probability suboptimality guarantees for the returned policy, covering both near-optimality certification and stopping at the certification limit. In the certifiable regime, under the stated assumptions,

the leading query-complexity term is governed by the same local quadratic characteristic, evaluated at a polylogarithmically reduced effective offline sample size. We additionally establish robustness to row-wise simulator misspecification measured in total variation.

Empirically, we evaluate the method on three synthetic environments (contextual bandit, multi-stage MDP, and an MDP with long-term consequences), and a semi-synthetic drug-response application. Compared with uniform and coverage-based allocation, our adaptive allocation certifies an $\varepsilon$-optimal policy with 3-22 times fewer collaborator queries in an efficient implementation and returns better policies at equal budgets.

Related Work

Worst-case sample complexity in tabular RL. A large body of work studies the worst-case sample complexity of learning a near-optimal policy in tabular MDPs. In offline RL, where the learner must rely entirely on a fixed dataset, guarantees of related work necessarily depend on how well the behavior policy covers the state–action pairs relevant to an optimal policy. A sequence of works has established sample complexity guarantees under various notions of coverage, including single-policy concentrability [24, 19].

At the other extreme, the generative-model setting allows the learner to request samples from arbitrary state–action pairs. Classical and more recent results characterize the minimax sample complexity of uniformly querying under such access [14, 1].

Instance-dependent policy identification. Our work is closely related to a growing literature on instance-dependent reinforcement learning, which seeks to exploit the specific structure of the underlying MDP to obtain substantially sharper sample-complexity guarantees than worst-case bounds. Rather than characterizing difficulty solely through global quantities such as the size of the state–action space or worst-case coverage coefficients, these results adapt to instance-specific features such as visitation probabilities, local variances, and suboptimality gaps. Hence the resulting guarantees can be much smaller than worst-case rates on favorable instances.

In offline tabular RL, Yin and Wang [2021] derived variance-sensitive instance-dependent guarantees while Wang et al. [2022] showed that, under suitable coverage conditions, the sample complexity can be improved when the minimum suboptimality gap is bounded away from zero.

A related line of work studies instance-dependent policy identification in online RL. For online tabular RL, Wagenmaker et al. [2022] developed instance-dependent PAC guarantees that adapt to local action gaps, thereby allowing substantially smaller sample complexity on MDPs in which suboptimal actions can be identified quickly. In contrast, Wagenmaker and Jamieson [2022] studied instance-dependent policy identification from a policy-level perspective, with complexity controlled by the suboptimal policy gap. More recently, Narang et al. [2024] showed that directly estimating value differences between candidate policies can lead to substantially sharper identification guarantees than estimating each policy value separately.

Instance dependence has also been studied under generative-model access. For discounted tabular MDPs, Zanette et al. [2019] derived structure-aware sample-complexity guarantees that adapt to local reward and transition variances as well as suboptimality action gaps, and can avoid an explicit polynomial dependence on the effective horizon in favorable instances.

Imitation learning and learning with expert queries. Several ideas underlying our framework also appear in the imitation-learning literature, particularly the selective use of expert supervision and the combination of expert information with other sources of data. Judah et al. [2014] studied active imitation learning, where the learner adaptively selects informative states at which to query the expert for an action, with the goal of reducing the number of expert labels. Liu et al. [2022] studied adversarial policy imitation with linear function approximation and, in their offline setting, combined expert demonstrations with an additional offline dataset collected by a prior policy. They showed that this supplementary dataset can improve the statistical efficiency of imitation learning under suitable coverage conditions. More closely related to our use of the collaborator, Amortila et al. [ [3] considered efficient learning with a small number of expert-action queries. In their setting, an expert query provides only the expert’s action and the learner implements it and interacts with the environment, and the goal is to learn a policy that performs at least nearly as well as the expert rather than to identify a near-optimal policy.

Adaptive querying and complementary information. Efficient information acquisition is also studied beyond RL. Ding et al. [ [11] adaptively select questions and respondents to infer latent group properties under budget constraints, while Liu et al. [ [20] allocate reasoning samples across questions for efficient test-time self-consistency. Related work also exploits complementary information: weak-strong model collaboration combines specialized knowledge with general reasoning [16]. Our setting instead concerns near-optimal policy identification from offline data and adaptively queried simulator observations. The collaborator’s fixed action policy restricts which uncertainties can be reduced, making both query efficiency and the attainability of the target accuracy central to the analysis.

Setup

In this section, we formalize the reinforcement learning setting we study.

Tabular MDPs. Following Li et al. [ [19], we consider a finite-horizon, time-inhomogeneous tabular MDP

where $\mathcal{S}$ and $\mathcal{A}$ are finite state and action spaces of sizes $|\mathcal{S}|$ and $|\mathcal{A}|$, $H$ is the horizon, and $\rho$ is the initial-state distribution. Write $[m] := \{1,\ldots,m\}$ for integers $m \ge1$. At stage $h \in[H]$, the transition kernel $\mathbf{P}_h$ assigns probability $P_h(s' \mid s,a)$ to next state $s'$ from $(s,a)$. The random reward vector $\mathbf{R}_h \in[0,1]^{|\mathcal{S}||\mathcal{A}|}$ has components $R_h(s,a)$ with means $r_h(s,a) := \mathbb{E}[R_h(s,a)]$. For every $(h,s,a)$, we assume that, conditional on $(S_h,A_h)=(s,a)$, the reward and next state are independent.

Preliminaries for analysis. A policy $\pi=\{\pi_h\}_{h=1}^{H}$ assigns action distributions $\pi_h(\cdot\mid s)\in\Delta_{\mathcal{A}}$, where $\Delta_{\mathcal{A}}$ is the probability simplex over $\mathcal{A}$. Let $\Pi_d$ denote the deterministic policy class, with $\pi_h(s)\in\mathcal{A}$ the selected action. Write $\mathbb{E}_{\mathcal{M}}^\pi$ and $\mathbb{P}_{\mathcal{M},\pi}$ for trajectory expectation and probability under $\pi$ in $\mathcal{M}$. Define the value and action-value functions by

with the terminal convention $V_{H+1}^\pi(s)=0$. The initial value is $V_1^\pi:=\mathbb{E}_{S_1\sim\rho}[V_1^\pi(S_1)]$. The state and state-action occupancies are

We further define the optimal value and the set of optimal deterministic policies by

The policy gap

defines the $\varepsilon$-optimal set

For any $\pi^\star\in\Pi^\star(\mathcal{M})$, define

This definition does not depend on the choice of $\pi^\star$. In particular, if the optimal deterministic policy is unique, then $\Delta_{\min}$ is the smallest strictly positive suboptimality gap, whereas if there are multiple optimal deterministic policies, then $\Delta_{\min}=0$.

Empirical estimation. We next construct the empirical MDP from a given dataset $\mathcal{D}$ and define corresponding policy values and visitation probabilities. For each $(h,s,a)$, suppose the dataset $\mathcal{D}$ contains $N_{hsa}$ reward-next-state observations $\{(R_{hsa}^{(i)},S_{hsa}^{\prime(i)})\}_{i=1}^{N_{hsa}}$, such that

For $N_{hsa}>0$, define

These estimates define the empirical MDP $\widehat{\mathcal{M}}^{\mathcal{D}}=(\mathcal{S},\mathcal{A},H,\rho,\widehat{P}^{\mathcal{D}},\widehat{r}^{\mathcal{D}})$. For any policy $\pi$, let $\widehat{V}^{\pi,\mathcal{D}}_h$ and $\widehat{Q}^{\pi,\mathcal{D}}_h$ denote its value and action-value functions under this MDP. Its state and state-action occupancies are