# Reward-Mixing MDPs with a Few Latent Contexts are Learnable

arXiv: https://arxiv.org/abs/2210.02594 · HTML: https://mlsanity.com/paper/2210.02594/

## Key facts

- Cited by 5 papers in the map (data through 14 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.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)

## Builds on (most cited first, 26 of 26)

- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Unifying Count-Based Exploration and Intrinsic Motivation](https://mlsanity.com/paper/1606.01868/) (2016)
- [Minimax Regret Bounds for Reinforcement Learning](https://mlsanity.com/paper/1703.05449/) (2017)
- [#Exploration: A Study of Count-Based Exploration for Deep Reinforcement Learning](https://mlsanity.com/paper/1611.04717/) (2016)
- [Contextual Markov Decision Processes](https://mlsanity.com/paper/1502.02259/) (2015)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [Online Clustering of Bandits](https://mlsanity.com/paper/1401.8257/) (2014)
- [When Is Partially Observable Reinforcement Learning Not Scary?](https://mlsanity.com/paper/2204.08967/) (2022)
- [Explore First, Exploit Next: The True Shape of Regret in Bandit Problems](https://mlsanity.com/paper/1602.07182/) (2016)
- [Adaptive Reward-Free Exploration](https://mlsanity.com/paper/2006.06294/) (2020)
- [RL for Latent MDPs: Regret Guarantees and a Lower Bound](https://mlsanity.com/paper/2102.04939/) (2021)
- [Sample Complexity of Multi-task Reinforcement Learning](https://mlsanity.com/paper/1309.6821/) (2013)
- [Reinforcement Learning of POMDPs using Spectral Methods](https://mlsanity.com/paper/1602.07764/) (2016)
- [Near-optimal Representation Learning for Linear Bandits and Linear RL](https://mlsanity.com/paper/2102.04132/) (2021)
- [PAC Reinforcement Learning for Predictive State Representations](https://mlsanity.com/paper/2207.05738/) (2022)
- [Provable Reinforcement Learning with a Short-Term Memory](https://mlsanity.com/paper/2202.03983/) (2022)
- [A PAC RL Algorithm for Episodic POMDPs](https://mlsanity.com/paper/1605.08062/) (2016)
- [Provably Efficient Reinforcement Learning in Partially Observable Dynamical Systems](https://mlsanity.com/paper/2206.12020/) (2022)
- [Learning in Observable POMDPs, without Computationally Intractable Oracles](https://mlsanity.com/paper/2206.03446/) (2022)
- [On the Minimax Optimality of the EM Algorithm for Learning Two-Component Mixed Linear Regression](https://mlsanity.com/paper/2006.02601/) (2020)
- [Reinforcement Learning from Partial Observation: Linear Function Approximation with Provable Sample Efficiency](https://mlsanity.com/paper/2204.09787/) (2022)
- [Beyond the Low-Degree Algorithm: Mixtures of Subcubes and Their Applications](https://mlsanity.com/paper/1803.06521/) (2018)
- [Reinforcement Learning in Reward-Mixing MDPs](https://mlsanity.com/paper/2110.03743/) (2021)
- [Source Identification for Mixtures of Product Distributions](https://mlsanity.com/paper/2012.14540/) (2020)
- [Learning Mixtures of Discrete Product Distributions using Spectral Decompositions](https://mlsanity.com/paper/1311.2972/) (2013)
- [Coordinated Attacks against Contextual Bandits: Fundamental Limits and Defense Mechanisms](https://mlsanity.com/paper/2201.12700/) (2022)

## Built on by (most cited first, 5 of 5)

- [RL in Latent MDPs is Tractable: Online Guarantees via Off-Policy Evaluation](https://mlsanity.com/paper/2406.01389/) (2024)
- [Test-Time Regret Minimization in Meta Reinforcement Learning](https://mlsanity.com/paper/2406.02282/) (2024)
- [Prospective Side Information for Latent MDPs](https://mlsanity.com/paper/2310.07596/) (2023)
- [A Classification View on Meta Learning Bandits](https://mlsanity.com/paper/2504.04505/) (2025)
- [On Computing Total Variation Distance Between Mixtures of Product Distributions](https://mlsanity.com/paper/2605.03839/) (2026)

## 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.