mlsanityData to 7 Oct 2026

Reward-Mixing MDPs with a Few Latent Contexts are Learnable

arXiv 2210.02594 ↗First posted 5 Oct 2022cs.LG cs.IT math.ITcited by 5

On the map: LLM alignment / RLHF / preference optimisation → Reward modeling and RLHF

Abstract (arXiv)

We consider episodic reinforcement learning in reward-mixing Markov decision processes (RMMDPs): at the beginning of every episode nature randomly picks a latent reward model among M candidates and an agent interacts with the MDP throughout the episode for H time steps. Our goal is to learn a near-optimal policy that nearly maximizes the H time-step cumulative rewards in such a model. Previous work established an upper bound for RMMDPs for M=2. In this work, we resolve several open questions remained for the RMMDP model. For an arbitrary M≥2, we provide a sample-efficient algorithm--textttEM2--that outputs an ε-optimal policy using tildeO left(ε-2 · Sd Ad · textttpoly(H, Z)d right) episodes, where S, A are the number of states and actions respectively, H is the time-horizon, Z is the support size of reward distributions and d=min(2M-1,H). Our technique is a higher-order extension of the method-of-moments based approach, nevertheless, the design and analysis of the algname algorithm requires several new ideas beyond existing techniques. We also provide a lower bound of (SA)^Omega(√M) / ε2 for a general instance of RMMDP, supporting that super-polynomial sample complexity in M is necessary.

Key facts. Cited by 5 papers in the map (data through 7 Oct 2026). Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation. Builds on 26 of its 48 references that are in the map. First posted on arXiv 5 Oct 2022.

Builds on

26 of 48 references.

Closest ancestors · 1

Same direction, newest first

  1. 2021Reinforcement Learning in Reward-Mixing MDPs7 Oct

Other references · 25

  1. 2013Playing Atari with Deep Reinforcement Learningcited by 3,211
  2. 2016Unifying Count-Based Exploration and Intrinsic Motivationcited by 632
  3. 2017Minimax Regret Bounds for Reinforcement Learningcited by 500
  4. 2016#Exploration: A Study of Count-Based Exploration for Deep Reinforcement Learningcited by 307
  5. 2015Contextual Markov Decision Processescited by 213
  6. 2019VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learningcited by 210
  7. 2022When Is Partially Observable Reinforcement Learning Not Scary?cited by 100
  8. 2014Online Clustering of Banditscited by 99
  9. 2016Explore First, Exploit Next: The True Shape of Regret in Bandit Problemscited by 78
  10. 2020Adaptive Reward-Free Explorationcited by 77
  11. 2021RL for Latent MDPs: Regret Guarantees and a Lower Boundcited by 69
  12. 2013Sample Complexity of Multi-task Reinforcement Learningcited by 58
13 more
  1. 2016Reinforcement Learning of POMDPs using Spectral Methodscited by 51
  2. 2021Near-optimal Representation Learning for Linear Bandits and Linear RLcited by 45
  3. 2022PAC Reinforcement Learning for Predictive State Representationscited by 42
  4. 2022Provable Reinforcement Learning with a Short-Term Memorycited by 40
  5. 2016A PAC RL Algorithm for Episodic POMDPscited by 37
  6. 2022Provably Efficient Reinforcement Learning in Partially Observable Dynamical Systemscited by 35
  7. 2022Learning in Observable POMDPs, without Computationally Intractable Oraclescited by 30
  8. 2020On the Minimax Optimality of the EM Algorithm for Learning Two-Component Mixed Linear Regressioncited by 24
  9. 2018Beyond the Low-Degree Algorithm: Mixtures of Subcubes and Their Applicationscited by 15
  10. 2022Reinforcement Learning from Partial Observation: Linear Function Approximation with Provable Sample Efficiencycited by 15
  11. 2020Source Identification for Mixtures of Product Distributionscited by 11
  12. 2013Learning Mixtures of Discrete Product Distributions using Spectral Decompositionscited by 7
  13. 2022Coordinated Attacks against Contextual Bandits: Fundamental Limits and Defense Mechanismscited by 4

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2023: 1; 2024: 2; 2025: 1; 2026: 1.

5 papers.

  1. 2024RL in Latent MDPs is Tractable: Online Guarantees via Off-Policy Evaluationcited by 5
  2. 2024Test-Time Regret Minimization in Meta Reinforcement Learningcited by 5
  3. 2023Prospective Side Information for Latent MDPscited by 4
  4. 2025A Classification View on Meta Learning Banditscited by 1
  5. 2026On Computing Total Variation Distance Between Mixtures of Product Distributionscited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2022Near-Optimal Regret Bounds for Multi-batch Reinforcement Learning
  2. 2021Policy Finetuning: Bridging Sample-Efficient Offline and Online Reinforcement Learning
  3. 2025Optimal Single-Policy Sample Complexity and Transient Coverage for Average-Reward Offline RL
  4. 2026Learning Kernel-Based MDPs from Episodic Preferential Feedback
  5. 2021Nearly Horizon-Free Offline Reinforcement Learning
  6. 2020Toward the Fundamental Limits of Imitation Learning
  7. 2025Generalization in Monitored Markov Decision Processes (Mon-MDPs)
  8. 2021Provably Breaking the Quadratic Error Compounding Barrier in Imitation Learning, Optimally