Source-linked AI summary
Markov Decision Processes with Continuous Side Information
Aditya Modi, Nan Jiang, Satinder Singh, Ambuj Tewari
TL;DR
The paper asks how reinforcement learning can share data across episodic MDPs when observed contexts identify changing dynamics, especially when data for any individual is limited. It develops PAC algorithms for smoothly varying CMDPs and for linear combinations of fixed MDPs, finding that smooth settings incur necessary exponential dependence on context dimension while the linear structure substantially improves PAC complexity. The analysis assumes known smoothness information for Cover-Rmax and does not seek the tightest possible bounds.
Problem
Learning personalized policies requires exploiting context to share relevant interaction data across many MDPs without assuming that all patients or users have identical dynamics.
Method
The paper proposes Cover-Rmax for smoothly context-varying MDPs and KWIK LR-Rmax, using KWIK linear regression, for contexts that form linear combinations of finitely many fixed MDPs.
Results
The smooth setting has a necessary exponential dependence on context dimension, while the linear structural assumption yields a significant PAC sample-complexity improvement.
Takeaways & Limitations
Observed side information supports near-optimal policy learning across a potentially infinite collection of MDPs, with substantially better PAC complexity available under linear structure.
Takeaways & Limitations
The Cover-Rmax PAC bounds are not intended to be tightest possible, and the analysis assumes the distance metric and smoothness constants Lp and Lr are known.
Abstract
from arXiv · showhide
We consider a reinforcement learning (RL) setting in which the agent interacts with a sequence of episodic MDPs. At the start of each episode the agent has access to some side-information or context that determines the dynamics of the MDP for that episode. Our setting is motivated by applications in healthcare where baseline measurements of a patient at the start of a treatment episode form the context that may provide information about how the patient might respond to treatment decisions. We propose algorithms for learning in such Contextual Markov Decision Processes (CMDPs) under an assumption that the unobserved MDP parameters vary smoothly with the observed context. We also give lower and upper PAC bounds under the smoothness assumption. Because our lower bound has an exponential dependence on the dimension, we consider a tractable linear setting where the context is used to create linear combinations of a finite set of MDPs. For the linear setting, we give a PAC learning algorithm based on KWIK learning techniques.
1. Introduction
The paper frames personalized sequential decision-making as a data-sharing trade-off and uses observed contexts to pool information across similar individuals. It studies CMDP sample complexity under smoothness and linear-structure assumptions.
- Personalized treatment requires balancing patient-specific policies against pooling data that may be irrelevant to the current patient.The motivation is healthcare, where limited per-patient data makes separate MDP models unreliable while full pooling can ignore patient differences.
- Observed side information can guide personalized policies and identify which individuals’ interaction data should be pooled.Examples include laboratory data and medical history in healthcare, user profiles in advertising, and student profiles in tutoring.
- The paper studies worst-case CMDP sample complexity in two settings: smoothly varying MDPs and linear combinations of finitely many fixed MDPs.It proposes Cover-Rmax with PAC bounds for the smooth setting and KWIK LR-Rmax with a PAC upper bound for the linear setting.
2. Contextual Markov Decision Process
A CMDP maps each observed context to an episodic MDP, with the agent selecting policies after observing contexts and then interacting for one episode. The analysis permits arbitrary context sequences and measures learning by the number of non-ϵ-optimal episodes.
- MDP setup: An MDP specifies states, actions, transitions, rewards, and an initial-state distribution, with policies acting over a fixed episodic horizon H.The value of a policy is compared with the optimal value achieved by an optimal policy.
- Relation to prior CMDPs: The contextual model extends prior CMDP work by assuming contexts are fully observed before interaction begins.Hallak et al. (2015) assumes latent contexts, creating differences in applications, assumptions, and results.
- Contextual MDP: A CMDP maps each context c to MDP parameters, while the paper assumes bounded contexts, rewards in [0,1], and a context-independent initial-state distribution.The context space is bounded in ℓ2 norm, and each episode begins by observing its context.
- Online protocol: At each episode, the agent observes the context, chooses a policy using prior experience, and interacts with the corresponding MDP.The policy may depend on the current context and previous episodes.
- PAC objective: The context sequence may be arbitrary or adversarial, and PAC analysis counts episodes in which the policy value is more than ϵ below optimal.The authors note that their Cover-Rmax PAC bounds are not designed to be tightest possible and could be combined with other PAC algorithms for improved dependence on S, A, and H.
3. Cover-Rmax
Cover-Rmax groups nearby contexts into finitely many cover balls, pools their data, and applies Rmax-style exploration under smoothness assumptions. Its upper bound scales with the context covering number, while a matching lower-bound construction shows this dependence is unavoidable in general.
- Cover-Rmax: The algorithm uses Rmax’s known-versus-unknown state-action distinction, optimistic rewards for unknown states, and balanced exploration to estimate contextual models.Counts and rewards are aggregated within cover balls, then used to construct estimated transition and reward models.
- Cover-Rmax: Cover-Rmax pools observations from nearby contexts within a finite cover, trading approximation bias against the number of context balls.The cover radius must be small enough to preserve smoothness-based accuracy, while pooled data avoids learning infinitely many MDPs separately.
- Cover-Rmax: The PAC guarantee bounds non-ϵ-optimal episodes using a cover radius r0 = min(ϵ/(8HLp), ϵ/(8Lr)) and a covering number N(C, r0).The proof adapts episodic Rmax analysis, requiring sufficiently many visits for accurate transition and reward estimates and applying a union bound across cover balls.
- Cover-Rmax: For d-dimensional Euclidean contexts, the covering-number dependence is O(1/r^d), producing exponential dependence on context dimension.The paper notes that this dependence reflects the difficulty of optimally learning under smoothness alone.
- 3.1 Lower Bound: The lower bound embeds independent hard episodic MDPs at separated context packing points, forcing separate learning and making linear dependence on the packing, and hence covering, number unavoidable.The construction satisfies smoothness with Lp = 1 and multiplies the single-MDP lower bound across packing points.
4. Contextual Linear Combination of MDPs
The linear setting models each contextual MDP as a combination of shared base MDPs, enabling global parameter sharing and KWIK-based exploration. This structure avoids the exponential dependence on context dimension that is unavoidable under smoothness alone.
- 4. Contextual Linear Combination of MDPs: The linear model represents each contextual MDP’s parameters as a combination of d base MDPs, with the context providing the combination coefficients.The context is restricted to the probability simplex so the resulting transition function remains valid.
- Model estimation: Data from one context can help estimate MDP parameters for other, potentially distant contexts because the base parameters are shared globally.This global structure replaces local pooling over context-space balls with linear regression over contexts.
- Model estimation: The method estimates transition vectors with regularized linear regression, using context vectors as features and next-state indicators as labels.The inverse empirical covariance determines whether predictions are accurate in the direction of a new context.
- 4.1 KWIK LR-Rmax: The vector-valued KWIK extension treats an entire transition distribution as known when its L1 prediction error is sufficiently small, avoiding separate accuracy tests for every next state.This yields tighter analysis than checking each scalar next-state probability independently.
- 4.1 KWIK LR-Rmax: KWIK LR-Rmax retains Rmax exploration while using KWIK predictions to mark state-action pairs known or unknown under the current context.Unknown predictions receive optimistic Rmax rewards, encouraging exploration and triggering updates when those pairs are observed.
- 4.1 KWIK LR-Rmax: The PAC analysis gives KWIK LR-Rmax a combined dependence on S and d of O(max{d4S, d2S3}), avoiding the exponential dependence on context dimension.The theorem bounds suboptimal episodes and the KWIK analysis bounds prediction failures and updates.
5. Related work
The paper relates observed-context CMDPs to latent-context transfer, side-information methods, metric-space RL, and KWIK-based linear learning. It emphasizes that observed contexts avoid identifying a latent MDP within a short episode, while its PAC guarantees differ from regret guarantees.
- Transfer in RL with latent contexts: Unlike latent-context transfer methods, observed side information avoids first identifying the underlying MDP, which may be impossible when an episode is only as long as the planning horizon.This distinction is motivated by limited interaction with individual patients, users, or students.
- RL in metric space: The smooth-CMDP approach pools observations across similar contexts rather than augmenting the state with context and applying metric-state PAC methods.The paper argues that state augmentation may increase sample and computational complexity unnecessarily.
- KWIK learning of linear hypothesis classes: The linear-combination method builds on KWIK-Rmax and KWIK linear regression to obtain PAC guarantees whose complexity depends on the linear-regression KWIK bound.It is presented as a parametric alternative yielding substantially improved PAC bounds.
- KWIK learning of linear hypothesis classes: Compared with a UCRL2 confidence-set approach, this work uses an independent method and provides PAC guarantees that are not directly comparable to regret bounds.The paper notes optimal dependence on A for PAC but suboptimal dependence on 1/ϵ relative to the competing regret setting.
6. Conclusion
The paper develops PAC learning with observed side information for large or infinite families of MDPs. Cover-Rmax handles smooth variation, while KWIK LR-Rmax improves sample efficiency under a linear-combination structure.
- 6. Conclusion: Cover-Rmax is a model-based PAC-exploration algorithm for MDPs that vary smoothly with observed side information.The setting targets near-optimal policies across a large or potentially infinite number of MDPs.
- 6. Conclusion: The lower-bound construction shows that smooth CMDPs require exponential dependence on context dimension for PAC learning in the worst case.The linear setting is introduced to obtain substantially improved guarantees under additional structure.
- 6. Conclusion: Under a linear combination of MDPs, KWIK LR-Rmax provides efficient exploration and a PAC analysis with improved dependence on context dimension.The paper identifies personalized recommendations, healthcare treatment policies, and tutoring systems as application possibilities.
A.1 Proof of Lemma 3.3
The proof analyzes episodic exploration using induced MDPs, simulation inequalities, and bounded visits to unknown states. It assumes rewards lie in [0,1] and counts complete episodes rather than every suboptimal action.
- Proof setup: The analysis adapts Kakade (2003) to episodic learning, removing a factor of H because complete episodes count as mistakes.The proof assumes rewards lie between 0 and 1.
- Induced MDP: The induced MDP keeps known-state dynamics and makes unknown states absorbing with reward 1.For states outside K, transitions remain at the same state and rewards equal 1.
- Simulation analysis: A simulation lemma relates transition and reward discrepancies between episodic MDPs to policy-value differences.The proof represents trajectory probabilities and expected average rewards under a fixed policy.
- Policy comparison: The induced-MDP inequalities and implicit explore-and-exploit corollary compare policies and values in the original and induced MDPs.The argument invokes Lemmas A.2 and A.3 to transfer approximation guarantees to the original MDP.
- Exploration bound: Balanced wandering bounds Rmax exploration by at most mSA visits to unknown states, plus H−1 delayed visits from episode-level updates.Episodes with escape probability above ϵ require enough exploration steps to exhaust this visit budget, with concentration supplied by Hoeffding’s inequality.
A.2 Proof of Theorem 3.2
The proof chooses a cover resolution and sample count sufficient to approximate transition and reward functions uniformly within each context ball. Concentration bounds and the simulation lemma then yield the desired value-function approximation.
- A.2 Proof of Theorem 3.2: The cover resolution and number of transitions m are selected to guarantee the required value-function approximation.The proof establishes this through the cover-approximation lemma.
- A.2 Proof of Theorem 3.2: Empirical averages over m transitions in each ball estimate transition probabilities for every context in that ball.The construction encodes observed next states as indicator vectors and averages them.
- A.2 Proof of Theorem 3.2: Hoeffding bounds control transition-estimation error simultaneously over state-action pairs with high probability.The proof obtains the bound with probability at least 1 − δ/2 before combining it with reward-error control.
- A.2 Proof of Theorem 3.2: Transition errors are reduced to ϵ/4H and reward errors to ϵ/4 for every context in the ball.These componentwise bounds provide the accuracy needed for the final value approximation.
- A.2 Proof of Theorem 3.2: The simulation lemma converts the transition and reward estimates into the desired approximation result.This is the final step of the proof.
B.1 Proof of Claim 3.6
The lower-bound proof places independently randomized hard MDPs at well-separated context packing points while preserving smoothness. Because observations at one point reveal nothing about another, learning requires substantial effort across many MDPs, producing the stated packing-number lower bound.
- B.1 Proof of Claim 3.6: The constructed contextual MDP is valid with transition smoothness Lp = 1 because its context-dependent probabilities satisfy the smoothness constraints.Rewards are constant across contexts, so reward smoothness holds for every Lr ≥ 0.
- B.1 Proof of Claim 3.6: Independent hard MDPs are assigned to context packing points, making the instances separately difficult to learn.The construction uses an 8ϵ′-packing and independent optimal-action choices at its points.
- B.1 Proof of Claim 3.6: Trajectories from other packing points provide no information about the instance selected at a given point.Under the product distribution over instances, contextual learning is equivalent to or harder than running a single-MDP learner at each packing point.
- B.1 Proof of Claim 3.6: The total number of mistakes is lower bounded by the aggregate difficulty across the packing points.The proof states this lower bound before converting the packing size into a covering-number expression.
- B.1 Proof of Claim 3.6: Setting |Z| = D(C, ϵ1) ≤ N(C, ϵ1) with ϵ1 = 8ϵ′ yields the stated lower bound.The final dependence is obtained by incorporating the packing or covering number into the known episodic-MDP lower bound.
Appendix C. Proof of the Theorem 4.1
The proof reduces vector-valued transition estimation to scalar linear regressions by projecting transition labels in every sign direction. Applying scalar KWIK guarantees and a union bound then establishes accurate transition predictions for the KWIK LR procedure.
- Appendix C. Proof of the Theorem 4.1: Vector-valued transition learning is reduced to scalar linear regression using a projection over sign vectors.For each f ∈ {−1, 1}^S, the transition label is projected to a scalar problem.
- Appendix C. Proof of the Theorem 4.1: The algorithm’s projected transition estimate equals the linear-regression prediction for the projected label because the regression is linear.Thus a KWIK guarantee for each projected problem transfers directly to the transition estimate.
- Appendix C. Proof of the Theorem 4.1: The proof bounds the ℓ1 error between estimated and true transition distributions for every context receiving a prediction.The ℓ1 norm is represented as a supremum over sign-vector projections.
- Appendix C. Proof of the Theorem 4.1: Scalar KWIK linear regression provides ϵ-accurate predictions with probability at least 1 − δ′ under bounded zero-mean noise and bounded coefficient norm.The theorem specifies the required algorithm parameter choice through α0 and gives a bound on the number of prediction states.
- Appendix C. Proof of the Theorem 4.1: Applying the scalar guarantee across projections and state-action contexts yields the transition-learning guarantee used by KWIK LR-Rmax.The proof sets δ′ = δ/(2S) and substitutes the relevant parameter values.