Projection-based Lyapunov method for fully heterogeneous weakly-coupled MDPs

Question 1: What is the relation between WCMDPs and factored MDP or constrained MDP?

Thanks for raising this interesting question. For the factored MDP and WCMDPs, while they both consist of some “loosely-coupled” components, their components are coupled in different ways: In the factored MDP, the components lie in a network, and the state transitions of nearby components are correlated. In contrast, the arms in WCMDPs have independent state transitions conditioned on the previous states and actions, but the joint actions of the arms are subject to a global constraint. It is unclear to us whether one problem class contains the other, or if there are any formal connections between these two classes.

WCMDPs is also different from the constrainted MDP: in the WCMDPs, the instantaneous costs in every time step are subject to constraints, whereas in the constrained MDP, the long-run average costs are subject to constraints. Nevertheless, the LP relaxation of the WCMDPs is equivalent to a constrained MDP, with (5b) being the long-run average constraints on each type of cost.

Question 2: How to go beyond unichain? Did the authors consider weakly-communicating models?

Going from unichain to weakly-communicating models (or more general) is non-trivial; even for the homogeneous case, asymptotically optimal policies are found only very recently by [22, 47]. On a high-level, their policies maintain the distribution of a subset of arms whose empirical state distribution is close to the optimal stationary distribution $\mu^*$, while trying to control the rest of the arms so that they can be gradually merged into the first subset. This idea might be useful for the fully heterogeneous setting, although an obvious obstacle is to remove the dependency on the empirical state distribution, and it will be an interesting future work to see if our projection idea can work here.

Question 3: l310: “let us assume …”: isn’t this implied by Assumption 1 ?

Thanks for bringing up this point. We would like to clarify that Assumption 1 only implies that each single-armed MDP mix to the steady state under the policy $\bar{\pi}_i^*$, whereas on Line 310 we assume that the joint state of all arms has a limiting distribution.

Question 4: l363, What do you mean by “strictly contractive”? Do you just mean contractive (Lipschitz-continuity with constant strictly smaller than 1)? Also in the proof of Lemma 3 (app E.2) I only found the proof of the Lipschitzness, not of the contraction property. Please clarify.

Thanks for pointing out this confusion. By “strictly contractive”, we are referring to the drift condition (C2) on page 8, i.e., $h(\mathbf{X}\_t P) \leq (1-C\_2) h(\mathbf{X}\_t) + O(\sqrt{N})$, and it is formally stated in the second bullet point of Lemma 3. We used the term “contractive” because $1-C\_2 < 1$, but we will change the wording to avoid this confusion.

Suggestion 2: What does “ID” mean (line 49 and onward)?

Here ID is a short hand for identity, which are simply the integer numbers 1, 2, 3, … N assigned to the arms. We will clarify this name in the revision.

We thank the reviewer for all the constructive suggestions.

Weakness: The main weakness I found in the paper is that the authors assume full knowledge of all the tabular MDPs. This assumption makes the problem significantly easier, as it allows the proposed policy to solve each MDP exactly ... I do not believe this is strictly a reinforcement learning paper, as it does not involve learning a policy through interaction with the environment(s).

We thank the reviewer for bringing up this point, as we also found the learning aspect of this problem interesting. That said, we would like to argue that the planning setting itself with known parameters is an interesting setting for the two reasons below, and it needs to be better understood before we move onto the learning setting.

First, the planning setting of WCMDPs and restless bandits has broad applications, especially in operations research, and there has been strong interest in developing better planning algorithms for more general settings. We’d like to refer to the following survey for a comprehensive overview, and the references in our paper for this continued effort.

José Niño-Mora. Markovian restless bandits and index policies: A review. *Mathematics*, 11(7), 2023. 

Furthermore, planning algorithms can play an instrumental role for learning a policy. For WCMDPs, if it is treated as a generic tabular MDP, then the sample complexity will be at least the state space size, which is exponential in N. A more tractable approach could be estimate model parameters first, and then apply a good planning algorithm with the estimated parameters. Such a model-based approach has been studied before for other RL problems. This approach is promising for WCMDPs with the ID policy and our results in this paper, but it is non-trivial to establish rigorous guarantees and sample complexity bounds. We will discuss the challenges in our response to Question 2.

Weakness (continued): Furthermore, I would have appreciated at least a toy simulation, to provide a practical illustration of the theoretical results.

We have performed simulation experiments on some random WCMDP instances. Below we display the simulation results on two representative instances. Each WCMDP instance is fully heterogeneous, whose arms are MDPs with independently generated parameters. For each arm, each row of the transition kernel is uniformly sampled from the probability simplex, each entry of the reward function and cost function is uniformly sampled from $[0, 1]$, and each budget parameter $\alpha_k$ is also uniformly sampled from $[0,1]$.

The first instance has a single cost constraint, with $|\mathbb{S}| = 10$, $|\mathbb{A}| = 4$,. We compare the performance of our policy with the performance of a policy name ERC (“Expected Reward and pull arms constrained by activation Cost” ) [46]. The ERC policy is designed for the typed heterogeneous WCMDPs with a single cost constraint, but we can apply it to fully heterogeneous WCMDPs by regarding each arm as having a different type. We evaluate the “optimality ratio” of the policies, which is defined as the ratio between the long-run average reward $R(\pi, \mathbf{S}_0)$ and the optimal value of the LP relaxation $R^{\text{rel}}_N$.

We simulate each of the two policies for $10^4$ time steps and $4$ replications for each $N\in \\{100,200,...1000\\}$. The result is displayed in the next table, where the confidence interval is calculated using the batch means method. We can see that the performance of these two policies are comparable, with the ID policy being slightly better.

The second instance has three cost constraints, with $|\mathbb{S}| = 10$, $|\mathbb{A}| = 2$,. We are not aware of any existing policies for the average-reward heterogeneous WCMDPs with multiple constraints, so we only evaluate the performance of our policy here. We simulate $10^4$ time steps and $4$ replications for each $N\in \\{100, 200, 300, ... 1000, 2000, 3000, ... 5000\\}$. Here, we consider a larger range of $N$ to more clearly demonstrate the asymptotic optimality of our policy. As we can see from the following table, as the number of arms $N$ gets large, the optimality ratio of our policy approaches $1$.

Due to the limited time of the rebuttal period, we have only performed simulations on a small number of instances. We could add a few more simulations and include the results in our final version if our paper gets accepted.

Question 1: Do you have any insights on how results extend to the case in which we consider continuous sub-MDPs?

We assume that the reviewer is referring to the continuous-space sub-MDPs. This is an interesting point, and we do think that our algorithm and analysis can possibly be extended to continuous spaces. To see this, observe that the real-time phase of the ID policy only needs to sample from the optimal single-armed policies $\bar{\pi}\_{i}^{\ast}$’s, which is possible even when the space is continuous. Moreover, our optimality gap does not explicitly depend on the sizes of state and action spaces, as can be seen from the expression of $C\_{\text{ID}}$ in second paragraph below Theorem 1. Therefore, if we can learn the optimal single-armed policies $\bar{\pi}\_{i}^{\ast}$ in the continuous sub-MDPs, our results can possibly carry through to continuous spaces.

Question 2: What are the main technical challenges in considering learning policies in unknown sub-MDPs? For instance, is would be possible to apply concepts by constrained reinforcement learning?

As mentioned at the beginning of this rebuttal, a reasonable approach for learning a policy for the N-armed WCMDP could be learning the model parameters and then applying the ID policy based on the learned parameters. To do this, we first learn the model parameters and construct the optimal single-armed policies for each sub-MDP. As the reviewer pointed out, this can be done through approaches for constrained RL. The main challenge in this approach is to handle how the errors in each sub-MDP affect the N-armed systems. Note that actions generated by single-armed policies may not be feasible in the N-armed system since they may not satisfy the hard constraints (this is where the ID policy steps in). Therefore, errors in each sub-MDP may have a non-linear effect and characterizing this effect is highly non-trivial.

Question 3: In the paragraph at lines 147--158, is it really necessary to indicate the policy when considering the vectors of MDPs states and actions?

We thank the reviewer for the pointing out this issue. We indicate the policy $\pi$ in superscripts mainly because $\pi$ is the decision variable of the optimization problem formulated in (3), and we want to emphasize the dependency on the decision variable. We can also omit $\pi$ in the paragraph at lines 147--158 if that is more natural.

Question 4: Is $R^*(N, \mathbf{S}\_0)$ only the optimal reward, or is it accounting for constraint satisfaction too?

It is the optimal reward; we consider the hard constraints and do not allow constraint violations.

Question 5: “In Lemma 1, what happens to constraint violation when employing the solution of the relaxed problem?”

Employing solutions of the relaxed problem would cause the constraint to be only satisfied in the long-run average sense, since the relaxed constraint (5b) is in terms of the long-run average state-action frequency $y(s,a)$. This is why we need Algorithms 1 and 2: we want to turn the solution of the relaxed problem into a policy that satisfies the hard constraint in (3b).

Weakness 1.1: I am not convinced by the problem setting's practicality. The average reward setting is not widely used since it is not easy to learn, and not as stable as the discounted reward setting. In this case, the authors should provide at least a synthetic toy case.

Thanks for bringing up this point. While learning in average-reward MDPs is harder than in discounted-reward MDPs, it is an active research area and increasingly efficient learning algorithms have been proposed. See, for example, the following paper for a review of the recent developments in learning average-reward MDPs:

Lee, Jongmin, Mario Bravo, and Roberto Cominetti. "Near-optimal sample complexity for MDPs via anchoring." *arXiv preprint arXiv:2502.04477* (2025).

For weakly-coupled MDPs, the policy that we propose can serve as a good baseline policy. We are currently working on leveraging such a near-optimal planning policy to develop learning algorithms with a lower sample complexity.

As for the experiments, we have performed simulations on some randomly-generated WCMDP instances, whose setup and data can be found in our reply to the Reviewer oixq. Due to the limited time of the rebuttal period, we have only performed simulations on a small number of instances. We could add a few more simulations and include the results in the final version if our paper gets accepted.

Weakness 1.2: Second, the hard instantaneous cost constraint is fixed for each time. It might be more practical to consider the cumulated cost for a period of time?

We agree that the cumulated cost is an interesting setting. That said, the hard instantaneous cost constraint is a natural model for allocating perishable resources, such the computing time of the servers, the availability of nursing personnels, the network bandwidth, etc.

Weakness 2: It will be better to provide more comparisons on the techniques with the discounted WCMDP literature, since the authors mentioned the optimality result is provided in that case.

Thanks for the suggestion. We have briefly discussed the difference between the existing results for discounted-reward WCMDPs and average-reward WCMDPs at the end of Appendix B. We could expand on that to comment more on the techniques.

At a high-level, the main challenge in the discounted setting is achieving transient optimality. To show that a certain policy achieves asymptotic optimality, one needs to show that the stochastic sample paths under the policy concentrate around an “optimal” deterministic path [See, e.g., 17, 52]. The concentration of sample paths usually follow from the concentration of per-step transitions, when the number of arms is large.

In contrast, in the averge-reward setting, the challenge mainly lies in understanding the long-term behavior of the system, and one needs to show that the steady-state distribution of the system state has some concentration properties. Such concentration properties do directly follow from the concentration of per-step transitions; intuitively, some global stability is needed to prevent the stochastic noise from accumulating over the infinite time. Establishing such global stability was a long-standing challenge resolved only recently [22,25,26,47], and a natural technique is to use Lyapunov functions.

Weakness 1: did not characterize the finite-time performance

Thanks for bringing up this point. In fact, our proofs of Theorem 1 can be easily adapted to give the following finite-time performance bound:

To prove this, we will only need to slightly modify Lemma 4 and 5. Note that the Lemma 5 actually gives a finite-time bound on the expectation of the Lyapunov function $V(\mathbf{X}\_t)$ (The equation above (82)), and Lemma 4 can be modified to show that:

Combining these two inequalities would give the finite-time performance bound that we claimed at the beginning.

We will add this result and its derivation to the paper.

Weakness 2: as $N\to\infty$, the assumption on the finiteness of mixing time may fall through.

Our theorem 1 currently assumes the existence of a constant-order upper bound $\tau$ on the mixing times of all arms (Assumption 1). While this assumption might sound strong, there are some straightforward ways to relax this assumption:

• First, if an $O(1/\sqrt{N})$ fraction of arms have unbounded mixing times, we can modify the algorithm to ignore these arms, which only causes $O(1/\sqrt{N})$ degradations in the optimality gap and does not affect the order of our bounds. Moreover, if we only aim for asymptotic optimality, i.e. $o(1)$ optimality gap, we can have at most $o(1)$ fraction of arms with unbounded mixing times.
• Second, while we mainly focus on the scaling of $N$, Theorem 1 remains true even if the mixing time $\tau$ scales up with $N$. In that case, Theorem 1 implies that the optimality gap of our policy is at most $O(\tau^4 / \sqrt{N})$ (see the second paragraph below Theorem 1), so asymptotically optimality is still achieved when $\tau = o(N^{1/8})$.

We could comment on these two relaxations of Assumption 1 in the paper.

Weakness 3: It is unclear as to what the approach should be if the LP relaxation has multiple optimal solutions.

We would like to point out that our policy and proofs do not rely on the uniqueness of the optimal solution to the LP relaxation. In the presence of multiple optimal solutions, any of them satisfying Assumption 1 can be used to define an instance of the ID policy, and Theorem 1 applies to all such instances. We can clarify this point in Section 3.

Weakness 4: $C\_{\text{ID}}$ seems to scale badly, making the result vacuous in large complex settings

We acknowledge that our analysis focuses on the scaling with $N$ and does not aim to optimize the dependence of the optimality gap bound on other parameters.

That said, we would like to highlight an interesting fact: $C\_{\text{ID}}$ does not explicitly depend on the sizes of the state or action spaces. This implies that our bound in Theorem 1 could still be informative for some WCMDP instances with large state-action spaces and many arms, provided that the relevant quantities (such as the mixing time bound $\tau$) remain moderate.

Question 1: In equation 8, is there a reason for choosing half the budget per cost constraint as the limit to determine the active constraints? Or is it an arbitrary choice?

Thanks for the question. This choice is arbitrary; any constant fraction strictly between $0$ and $1$ would yield the same $O(1/\sqrt{N})$ optimality gap. We will clarify this point in our revision.

Question 2: In equation 7, is there any intuition beyond setting the policy to be uniform in transient states?

Thanks for the question. To clarify, there is no unique “correct” way to define the single‑armed policy $\bar{\pi}\_{i}^{\ast}$ in transient states. For the purpose of proving Theorem 1, any definition of the transient behavior of $\bar{\pi}\_{i}^{\ast}$ suffices, provided it ensures the induced Markov chain is an aperiodic unichain with mixing time bounded by $\tau$ (Assumption 1). However, the specific choice may affect the required strength of Assumption 1.

The rationale behind our definition of $\bar{\pi}\_{i}^{\ast}$ is that it maximizes the "connectivity" of the induced chain. Specifically, if any definition of $\bar{\pi}\_{i}^{\ast}$ in transient states induces a unichain for arm $i$, then this uniform‑action definition also induces a unichain, since it assigns positive probability to every action.

Question 3: Does line 259 imply that as the arms with lower indices approach their stationary values, the cost constraints are satisfied in a sense that lets more arms follow their LP-optimal policies?

Your understanding is accurate. We will make this point more explicitly in our revision.

Question 4: It will help to include the intuition behind Algorithm 1 in the main body of the paper.

We thank the reviewer for the suggestions. We have provided some intuitions on Lines 263—269, but we could also move up some content in Appendix C to the main body.

Question 5: Is there any intuition behind choosing the cost and reward vectors as the feature vectors?

We design the feature vectors to preserve sufficient information of the system state $\mathbf{X}\_t$ so that, after projection, we can still determine (i) whether a subset of arms can follow the optimal single-armed policies, and (ii) the instantaneous reward of the subset of arms if they follow the optimal single-armed policies. The first goal requires us to include the cost vector so that we can check if the constraints are satisfied; the latter requires us to include the reward vector.

Question 6: Is the regret quantified with respect to the best stationary optimal policies or best optimal policies in a wider class of policies (as initially pointed at)?

Thanks for pointing out this confusion. We would like to clarify that these two ways of quantifying the regret are equivalent in our setting. Note that the WCMDP problem is an MDP with finite state and action spaces, so, by Theorem 9.1.8 of [35], there exists a stationary Markov policy that is optimal within the class of history-dependent randomized policies (the widest possible class). We will revise Lines 183-185 to clarify this.

Minor comment: In the beginning the authors mention that lim inf and lim sup of the average cost are the same in the finite state action setting. However, this is true subject to certain constraints on the Markov chain induced by the policy. Unichain and aperiodicity are sufficient (as assumed later) but this assumption needs to be made clearer earlier on.

We would like to clarify that we do not need the additional assumptions that the reviewer mentions. There are about two places where we mention that lim inf equals lim sup, which we clarify one-by-one:

• In the paragraph below (3), we are only arguing that the optimal limsup average reward equals the optimal liminf average reward, which always hold in the finite state-action setting, by Proposition 9.1.6 of [35] cited in this paragraph.
• At the top of Page 5, we argue that the limsup average reward equals the liminf average reward for a finite-state Markov chain. Note that we fix the initial state $\mathbf{S}\_0$ and consider the long-run average of the expectation, so we do not really need unichain or aperiodicity.

Dear Reviewer,

Thanks again for reviewing our paper and providing feedback. As the Author-Reviewer discussion phase is concluding soon, we would like to confirm whether we have fully addressed the questions you raised.

--Authors