PaperlayerView exact revision on arXiv

Source-linked AI summary

Provably efficient RL with Rich Observations via Latent State Decoding

Simon S. Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudík, John Langford

arXiv:1901.09018v3cs.LGstat.ML

TL;DR

The paper studies how to explore episodic RL environments with rich observations generated from a small latent state space, where hidden states complicate efficient exploration. It learns latent-state decoding through regression and clustering under identifiability assumptions, then uses the decoding to construct exploration policies with theoretical guarantees and strong empirical performance.

  • Problem

    Exploration in episodic RL with rich observations has limited theoretical sample-efficiency guarantees, especially when latent states are hidden and direct observation-space exploration is impractical.

  • Method

    The method inductively learns a mapping from rich observations to latent states using backward-probability embeddings, least-squares regression, and clustering, then constructs exploration policies.

  • Results

    The algorithms achieve sample complexity polynomial in latent-state, horizon, and decoding-class parameters without explicit observation-space dependence, and substantially outperform naïve exploration empirically.

  • Takeaways & Limitations

    Explicit latent-state recovery makes rich-observation exploration tractable under identifiability assumptions and supports policy covers for reaching latent states.

  • Takeaways & Limitations

    The guarantees depend inversely on the separability margin γ and the worst-case reaching probability µmin, and the formalism omits rewards from its central exploration task.

Abstract

from arXiv · show

We study the exploration problem in episodic MDPs with rich observations generated from a small number of latent states. Under certain identifiability assumptions, we demonstrate how to estimate a mapping from the observations to latent states inductively through a sequence of regression and clustering steps -- where previously decoded latent states provide labels for later regression problems -- and use it to construct good exploration policies. We provide finite-sample guarantees on the quality of the learned state decoding function and exploration policies, and complement our theory with an empirical evaluation on a class of hard exploration problems. Our method exponentially improves over $Q$-learning with naïve exploration, even when $Q$-learning has cheating access to latent states.

1. Introduction

The paper addresses exploration in episodic RL with rich observations, where sample efficiency is theoretically unclear and direct observation-based exploration is impractical. It learns latent-state decoding and exploration policies under structural assumptions, achieving observation-space-independent guarantees and strong empirical performance.

  • Rich-observation RL exploration has limited theoretical understanding, despite empirical algorithms designed for images and texts.
  • Treating observations as unique states makes tabular-style exploration unsuitable because interaction requirements depend polynomially on the observation space.
  • The method explicitly recovers a decoding function that maps rich observations to latent states, reducing the problem to tractable latent-state exploration when learned perfectly.
  • Under identifiability assumptions, learning requires samples polynomial in latent states, horizon, and decoding-class complexity, with no explicit observation-space dependence.
  • The algorithm is easy to implement and substantially outperforms naïve exploration empirically, including against baselines with cheating access to latent states.
  • The approach uses backward conditional probability vectors as latent-state representations and learns decoding through conditional-probability estimation and least-squares regression.

2. Setting and Task Definition

The paper formalizes rich observations through block Markov decision processes, where observable contexts are generated by unobservable latent states. Its exploration objective is an approximately optimal policy cover that reaches every latent state without scaling explicitly with the context-space size.

  • Block Markov Decision Process: A BMDP has finite unobservable latent states, finite actions, and a possibly infinite observable context space connected by transition and emission probabilities.
  • Block Markov Decision Process: The block structure makes each context identify its generating state, while the large context space motivates learning dynamics in the smaller latent space.
  • Block Markov Decision Process: The model partitions states and contexts by observable episode level, with at most M latent states per level and no explicit dependence on |X| sought.
  • Exploration Task: Exploration seeks policies that reach each latent state, with an h-step policy mapping observed contexts up to step h to actions.
  • Exploration Task: An ϵ-policy cover contains policies that reach every state within ϵ of its maximum reaching probability.
  • Exploration Task: A small policy cover enables data collection at sufficiently reachable states and supports obtaining a policy that maximizes any reward.

3. Embedding Approach

The embedding approach learns context and state representations in a shared low-dimensional space, using separable latent transitions to decode hidden states. Backward probability vectors provide the representations, which can be estimated by least-squares regression and separated by clustering.

  • Motivation: The central difficulty is decoding latent states without observing them, so the method assumes separability of latent transition probabilities rather than restrictive context-emission structure.
  • Embeddings and Function Approximation: Context and state embeddings map into a shared simplex so contexts are close to the embedding of their generating state.
  • Embeddings and Function Approximation: The context embedding is learned from a function class G whose inductive biases support generalization across rich contexts and control sample complexity.
  • Algorithmic Construction: Least-squares regression estimates context embeddings, after which clustering groups contexts generated by the same latent state because their embeddings are close.
  • Backward Probability Vectors: Backward probability vectors collect predecessor-action probabilities for each successor state and serve as the state embeddings approximated by the algorithm.
  • Backward Probability Vectors: γ-separability requires distinct successor states’ backward probability vectors under uniform predecessor-action sampling to differ by at least γ in ℓ1 distance.
  • Separability: Deterministic latent transitions automatically satisfy the assumption with γ = 2, while the framework also covers a substantially larger class of γ-separable models.

4. Algorithm for Separable BMDPs

Algorithm 1 learns latent-state representations and an exploration policy cover inductively for separable BMDPs. Its regression, clustering, and dynamic-programming steps yield finite-sample guarantees, with stronger exact-reachability guarantees for deterministic transitions.

  • Algorithm 1: Algorithm 1 constructs discovered latent states, a decoding function, estimated transitions, and level-wise policies inductively.At each level, previously learned objects support regression, clustering, transition estimation, and policy construction for the next level.
  • Regression step: Regression learns context embeddings from trajectories generated by a policy mixture and decoded predecessor states.The embeddings are estimated by empirical risk minimization using samples of decoded prior states, actions, and subsequent observations.
  • Clustering step: Clustering groups contexts whose learned embeddings are close, defining state centers and a decoder that assigns each context to its nearest center.The threshold-based clustering procedure produces discovered states and an embedding map, while the decoder uses nearest-center assignment.
  • Dynamic programming: Dynamic programming estimates latent transitions and constructs policies that approximately maximize the probability of reaching each target state.The resulting observation policies compose decoded-state policies with the learned decoding functions.
  • Guarantees: With probability at least 1 −δ, Algorithm 1 returns an ϵ-policy cover of S with size at most MH, under suitable accuracy and sample-complexity conditions.The guarantee depends on M, K, H, 1/ϵ, the separability margin γ, and the worst-case reaching probability µmin, while suppressing logarithmic factors.
  • Deterministic specialization: For deterministic BMDPs, the method obtains an ϵ-policy cover with ϵ = 0, and each learned policy reaches its corresponding latent state with probability one.Deterministic transitions also permit fixed action sequences and remove the need for dynamic programming when constructing the cover.

5. Experiments

The experiments evaluate PCID on hard combination-lock environments with rich observations, comparing it against tabular methods that directly access latent states. PCID substantially improves exploration scaling over naïve Q-learning, while remaining sensitive to violated assumptions and under-estimated hyperparameters.

  • Experimental setup: The six-environment evaluation compares PCID on rich observations with ORACLEQ and QLEARNING, which directly access latent states.Some environments satisfy BMDP assumptions and others test robustness when those assumptions fail.
  • Experimental setup: The combination-lock environments have H levels, three states per level, and four actions, with rewards available only from selected latent states.The deterministic variant contains 2^H rewarding paths among 4^H total paths, making random exploration exponentially difficult.
  • Experimental setup: Time-to-solve is the first episode whose running-average reward reaches 0.25 = 0.5V⋆, summarized by medians over 25 replicates.The experiments run for up to 100K episodes and report 90th and 10th percentile error bars.
  • Results: QLEARNING cannot solve problems with H ≥15 within 100K episodes, while its linear performance curve reveals exponential sample complexity.An extended run solved H = 15 in 170K episodes.
  • Results: In Lock-Bernoulli, PCID is roughly 5× worse than ORACLEQ but has similar behavior and exponentially improves over ε-greedy QLEARNING.PCID shows near-linear scaling with H, even better than predicted by the theory.
  • Results: PCID becomes less robust as Lock-Gaussian noise and stochasticity increase, although its sampling-complexity scaling remains favorable relative to QLEARNING.At σ = 0.3, PCID fails to solve stochastic H = 40 with linear functions and H = 30 with neural networks.
  • Sensitivity analysis: Under-estimating k or n causes failure, whereas over-estimating both parameters is relatively robust.Too few clusters can miss latent states, while too few trajectories can leave regression problems under-sampled.
  • Summary: PCID scales favorably across linear and nonlinear rich-observation problems, dramatically outperforming ε-greedy QLEARNING and remaining roughly 5× slower than ORACLEQ.The experiments explicitly do not measure regret.

B. Incorporating Rewards in BMDPs

The paper distinguishes reward models in BMDPs according to whether rewards depend only on latent states or also on contexts. Its policy cover supports direct policy learning in the first case and exploration datasets for off-policy methods in the second.

  • Latent-state rewards: When rewards depend only on latent states, near-optimal policies or value functions can be built using the paper’s policy-cover methods.This preserves the interpretation of the latent state as sufficient for an optimal policy or value function.
  • Context-dependent rewards: When rewards depend on contexts, latent states alone cannot construct optimal policies or value functions, but the policy cover still provides state-action coverage.The resulting dataset can be used by approximate dynamic programming or policy-learning methods.
  • Guarantee scope: The paper’s approach naturally targets PAC-style guarantees rather than regret minimization.Despite this mismatch, the empirical evaluation reports good regret performance on challenging benchmarks.

C.1. Implementation Details

The implementation uses regression or neural-network oracles, K-means clustering, and specified data-reuse and optimization procedures. Additional figures compare time-to-solve across observation settings and latent-transition conditions.

  • Data collection: The implementation sets Ng = Nφ = Np = n, collects n uniformly sampled trajectories, and reuses that data instead of collecting another n samples.These choices implement the corresponding data-collection steps of Algorithm 1.
  • Oracle implementation: Linear representations use independent unregularized least-squares regressions for each output coordinate, solved with a pseudoinverse and no intercept.For vector-valued predictions, regression is performed separately on each coordinate.
  • Oracle implementation: Neural-network oracles use a sigmoid hidden layer, square loss, AdaGrad, and a maximum of 5000 optimization iterations.The supplied implementation description specifies the weight dimensions, learning-rate multiplier, convergence check, and stopping limit.
  • Additional visualizations: Figure 3 plots time-to-solve against problem difficulty on a log-log scale for two observation processes and function classes.The figure description states that PCID and ORACLEQ curves show linear scaling with difficulty.
  • Additional visualizations: Figure 4 compares time-to-find across Lock-Bernoulli, Lock-Gaussian with linear functions, and Lock-Gaussian with neural networks under deterministic and stochastic transitions.The stochastic condition uses switching probability 0.1; ORACLEQ and QLEARNING operate directly on hidden states.
  • Clustering: PCID uses K-means clustering with model selection that decreases k when a cluster has fewer than 30 points.The implementation starts from the input k and searches for clusters meeting the minimum-size criterion.

C.2. Additional Results

Additional experiments examine scaling, exploration-specific time-to-find, and hyperparameter sensitivity. PCID displays near-linear scaling in the tested settings, while time-to-find separates discovery from later exploitation.

  • Scaling: 1.051: PCID’s log-log time-to-solve slope on Lock-Bernoulli is the same in deterministic and stochastic settings.Because d = H in these experiments, the authors report this linear dependence as substantially better than their theoretical prediction.
  • Scaling: 1.065 and 1.013: ORACLEQ’s deterministic and stochastic log-log slopes, respectively, verify near-linear scaling with H.The slopes correspond to the exponent on H in ORACLEQ’s sample complexity.
  • Exploration metric: Time-to-find records the first episode with non-zero total reward, making it more closely aligned with solving the exploration problem than time-to-solve.Time-to-solve additionally requires exploration and exploitation.
  • Exploration results: With neural-network representations, PCID almost always finds the goal within 100K episodes even when it does not accumulate high reward.The authors attribute this pattern either to exploitation failure, which is outside the paper’s focus, or to needing a few more episodes.
  • Sensitivity analysis: k = 3 and n = 300 are the best parameter choice shown for stochastic Lock-Bernoulli with H = 20.Underestimating either parameter causes failure, whereas overestimating both is described as robust.

C.4. Reproducibility Checklist

The reproducibility checklist describes the environments, randomization protocol, hyperparameter searches, repeated evaluations, and reported metrics. No dataset collection, train/test split, or data exclusion was used.

  • Data and environments: No dataset was collected; experiments used environments implemented through the OpenAI Gym API.The environment source code was included with the submission and intended for public release.
  • Protocol: The evaluation used different random seeds for development and final experiments instead of a train/validation/test split.The performance metrics were treated as regret-like quantities that require no split.
  • Data handling: No data was excluded from the experiments.
  • Hyperparameters: PCID used K-means parameter 3 and searched n from 100 through 1000 in increments of 100, while baseline hyperparameters were separately tuned.The checklist lists learning-rate, confidence, and exploration-parameter grids for ORACLEQ and QLEARNING.
  • Evaluation: Each experiment used 25 replicates, with all algorithms run for 100K episodes under varied environment featurizations and stochasticity.Random seeds changed for both the environment and algorithm in every replicate.
  • Reported metrics: Reported metrics were median time-to-solve and time-to-find, with 90th- and 10th-percentile error bars over replicates.Time-to-solve uses a running-average reward threshold, whereas time-to-find uses the first non-zero reward episode.

D. Proofs for Deterministic BMDPs

The deterministic-BMDP proofs analyze regression error and inductively establish correct latent-state decoding and policy coverage. These guarantees yield an exact policy cover and an approximately optimal policy with high probability.

  • Proof strategy: The proofs begin with deterministic BMDPs because their arguments are simpler and reused in the stochastic analysis.
  • Regression analysis: Theorem D.1 bounds the learned regression function using empirical-risk-minimization analysis and Bernstein’s inequality.The proof defines population and empirical risks under uniformly sampled state-action pairs and emitted observations.
  • Main guarantee: With probability at least 1 − δ, Algorithm 4 returns an ε-policy cover with ε = 0 when τ = 0.01, Ng = Õ(M^2K^2 log |G|), and Nφ = Õ(MK).
  • Inductive invariants: The induction maintains bijection between decoded and true states and zero-scale coverage of each earlier level.The base case follows from the fixed starting state, and the induction proceeds level by level.
  • Clustering argument: At each level, clustering rejects embeddings near an already represented successor and adds embeddings separated from all existing clusters.The resulting decoded and true state sets are bijective, with one stored policy path reaching each state.
  • Policy guarantee: With probability at least 1 − δ, Algorithm 4 can find an ε-suboptimal policy using Õ trajectories in a deterministic BMDP.The corollary follows by estimating rewards for state-action pairs and applying concentration and union bounds.

E. Proof of Theorem 3.1

The proof shows that least-squares regression recovers the backward conditional probability representation of each latent state. This establishes the stated form of every regression minimizer.

  • Every minimizer gh satisfies gh(x′) = bν(s′) for all observations emitted by latent state s′.This is the theorem’s conclusion under Assumption 3.1.
  • The conditional mean of the previous state-action pair minimizes the least-squares loss for each next latent state.The proof identifies bν(s′) as this conditional mean and invokes its optimality under squared loss.
  • The tabular embedding φ(s′) = bν(s′) minimizes unconditional squared loss under any distribution over next latent states.Because φ is tabular, the conditional regression objective can be related to the unconditional objective.

F. Justification of Assumption 3.2 and Dependency on µmin

The section establishes that separability of backward probabilities is necessary for state identification and quantifies how nonuniform visitation affects the resulting margin. It also illustrates why dependence on the minimum visitation probability cannot generally be removed.

  • When γ = 0, the uniform backward representation cannot differentiate the corresponding latent states.The section states that no alternative distribution can restore differentiation in this case.
  • If the uniform backward probabilities of two latent states coincide, no strictly positive previous-state distribution can distinguish them.The theorem derives this implication using Bayes’ rule and diagonal reweighting of the backward probabilities.
  • If the uniform backward probabilities differ, every distribution assigning positive probability to each previous state-action pair preserves state distinguishability.The theorem gives the converse condition for identification using distributions over previous state-action pairs.
  • A minimum visitation condition ν(s,a) ≥ τ converts uniform separability γ into an ℓ1 margin of at least τγ/2.This is the quantitative conclusion of Lemma F.1 under Assumption 3.2.
  • The inverse dependence on µmin is unavoidable because an exponentially rare predecessor state cannot be reliably used for exploration.The example sets the visitation probability of s2 to exp(−1/ϵ), making differentiation through that state infeasible.

G. Proof of Theorem 4.1 and Claim 4.1

The proof proceeds inductively over levels, showing that learned states correspond bijectively to true states, decoding is accurate, and transition estimates support exploration. These properties yield an approximate policy cover with high probability.

  • Inductive proof: The induction maintains three properties: learned states correspond to true states, decoding is accurate, and estimated transitions are approximately correct.The proof decomposes Claim 4.1 into conditions established at each level.
  • Exploration guarantee: Each target state has a learned-state policy reaching its true counterpart with probability at least µ(s) − 2Hϵf − 2Hϵp.This is the key policy-reaching guarantee used to establish the policy-cover condition.
  • Exploration guarantee: With ϵf and ϵp sufficiently small, the algorithm outputs an ϵ-policy cover with high probability.The proof explicitly plugs the error bounds into the policy-reaching lemma to complete Theorem 4.1.
  • Inductive proof: The decoding and embedding estimates make same-state contexts close and different-state contexts well-separated, enabling clustering by latent state.Lemma G.2 supplies separation and regression accuracy, while the proof interprets these as clusterability conditions.
  • Inductive proof: Clustering supplies a bijection between learned and true states and preserves small decoding error.The proof constructs the map αh and establishes its correctness through the clustering guarantees.

G.1. Proof details for Theorem 4.1 and Claim 4.1

The proof compares true-state and learned-state policies through their decoding events and transition estimates. It then shows that inductive exploration provides sufficient coverage for regression, separation, clustering, and policy construction.

  • Policy construction: The learned and true policy analyses are linked by the one-to-one correspondence between learned and true states.The proof transfers policies through α and α^-1 while accounting for decoding events.
  • Policy construction: Dynamic programming uses the learned transition model to construct policies that approximately maximize reaching probabilities for specific learned states.These policies act on decoded latent states and are composed with the learned decoding functions.
  • Exploration coverage: The exploration mixture reaches every predecessor state-action pair with probability at least µ(s)/(2MK), providing data for transition estimation.Uniformly mixing the reaching policies and actions yields the stated coverage lower bound.
  • Representation learning: Backward-probability separation is inherited from the exploration distribution, while regression estimates the representations accurately with high probability.The proof applies the margin lemma to ηh and then establishes the corresponding guarantees for ˆgh.
  • Clustering: The clustering sample size ensures every latent state contributes a representative point and all sampled contexts are decoded correctly with high probability.The proof uses coverage, a union bound, and the decoding-error condition to establish these sampling events.
Loading 1901.09018v3…