Reward-Mixing MDPs with a Few Latent Contexts are Learnable

TL;DR

This paper introduces an efficient algorithm for learning near-optimal policies in reward-mixing MDPs with multiple latent reward models, resolving open questions for arbitrary M and establishing fundamental complexity bounds.

Contribution

It extends previous work to arbitrary M, providing a sample-efficient algorithm with theoretical guarantees and new techniques for higher-order moments in RMMDPs.

Findings

The exttt{EM}^2 algorithm achieves near-optimal policy learning with polynomial sample complexity.

A lower bound shows super-polynomial sample complexity is necessary in M.

The method generalizes the method-of-moments approach to complex reward-mixing scenarios.

Abstract

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\ge2$, we provide a sample-efficient algorithm--$\texttt{EM}^2$--that outputs an $\epsilon$-optimal policy using $\tilde{O} \left(\epsilon^{-2} \cdot S^d A^d \cdot \texttt{poly}(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(\sqrt{M})} / \epsilon^{2}$ for a general instance of RMMDP, supporting that super-polynomial sample complexity in $M$ is necessary.