Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning

Dipendra Misra; Mikael Henaff; Akshay Krishnamurthy; John Langford

Source-linked AI summary

Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning

Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, John Langford

arXiv:1911.05815v1Nov 13, 2019cs.LGstat.ML

TL;DR

Exploration in rich-observation reinforcement learning is difficult when learning a compact state representation and exploring effectively are interdependent. HOMER learns kinematic inseparability abstractions and uses them for strategic exploration, proving polynomial trajectory complexity independent of observation-space size and demonstrating tractable performance on a difficult exploration problem.

Problem
Provably effective exploration remains elusive when learning a compact state decoder and exploring effectively are interdependent in rich-observation environments.
Method
HOMER learns kinematic inseparability abstractions from raw observations through supervised-learning reductions and uses them to drive reward-free exploration and policy optimization.
Results
HOMER learns a policy cover and near-optimal policies for Block MDPs with polynomially many trajectories and no explicit dependence on observation-space size, while solving a difficult exploration problem empirically.
Takeaways & Limitations
Kinematic inseparability provides a state abstraction that supports provably efficient exploration and reward optimization in rich-observation Block MDPs under reachability and function-class assumptions.
Takeaways & Limitations
Existing abstraction approaches based on predicting previous actions or abstract states can fail to produce a policy cover in general Block MDPs.

Abstract

from arXiv · show

We present an algorithm, HOMER, for exploration and reinforcement learning in rich observation environments that are summarizable by an unknown latent state space. The algorithm interleaves representation learning to identify a new notion of kinematic state abstraction with strategic exploration to reach new states using the learned abstraction. The algorithm provably explores the environment with sample complexity scaling polynomially in the number of latent states and the time horizon, and, crucially, with no dependence on the size of the observation space, which could be infinitely large. This exploration guarantee further enables sample-efficient global policy optimization for any reward function. On the computational side, we show that the algorithm can be implemented efficiently whenever certain supervised learning problems are tractable. Empirically, we evaluate HOMER on a challenging exploration problem, where we show that the algorithm is exponentially more sample efficient than standard reinforcement learning baselines.

1 Introduction

HOMER addresses the coupled challenge of representation learning and exploration in Block MDPs with rich observations by learning a kinematic state abstraction and using it to guide exploration. It achieves polynomial sample complexity independent of observation-space size and performs well on an extremely sparse-reward problem.

Motivation: Rich-observation exploration is statistically intractable in general, while learning a compact state decoder and tractably exploring are mutually dependent challenges.The paper motivates structured problem classes where sample-efficient algorithms may still be possible.
Problem and abstraction: HOMER targets Block MDPs, where rich observations are generated from a small number of unobserved latent states, and learns a reward-free state abstraction called kinematic inseparability.Kinematic inseparability aggregates observations with the same forward and backward dynamics.
Algorithm: HOMER interleaves exploration with abstraction refinement, training policies to visit each abstract state and forming a policy cover through a contextual-bandit reduction.The method uses dynamic programming and synthetic rewards that incentivize reaching an abstract state.
Guarantees: HOMER learns to visit every latent state and optimize any reward function with trajectories polynomial in latent states, actions, horizon, and function-class complexity, without explicit dependence on observation-space size.The guarantees assume reachable latent states and rely on two function classes used by the algorithm.
Empirical evaluation: 10−100 is the random-search probability of encountering an optimal reward in the empirical googal-sparse task, where HOMER recovers the abstraction and consistently finds a near-optimal policy.HOMER outperforms baselines using naive exploration and more sophisticated exploration bonuses.

2 Preliminaries

The paper formalizes episodic Block MDPs with finite hidden states, potentially infinite observations, and disjoint state-conditioned emissions. It defines reachability, homing policies, policy covers, function classes, and computational oracles as the foundations for learning and exploration.

Block MDP model: Block MDPs have potentially infinite observations but finite hidden states, finite actions, and horizon H, with agents choosing actions from observations without observing latent states.The model includes an initial-state distribution, transitions, emissions, and rewards over trajectories.
Environment assumptions: Disjoint emission supports remove partial-observability effects by defining an unknown inverse mapping from each observation to its latent state.This assumption motivates the Block MDP structure and is described as natural for visual grid-world scenarios.
Environment assumptions: The analysis assumes states are reachable with reasonable probability, with sample complexity scaling polynomially in the inverse minimum visitation probability, η^-1.Unlike prior work, it does not require additional identifiability or determinism assumptions.
Reward-free learning: A policy cover is a finite set of non-stationary policies that visits every state with probability at least αη(s), approximating the role of state-specific homing policies.Constructing such a cover supports collecting data throughout the state space and subsequently optimizing arbitrary rewards.
Function classes and computational oracles: Because observations may never repeat, learning uses function approximation through a policy class Π and regression classes built from abstract-state maps and tabular regressors.Computational oracles provide optimization primitives for these classes without statistical benefit, abstracting away search costs.

3 Kinematic Inseparability State Abstraction

Kinematic inseparability is a state abstraction that can be learned through supervised learning and supports reward-free exploration. Its backward variant suffices for exploration, while combining backward and forward dynamics yields the full abstraction and partitions observations into abstract states.

Kinematic inseparability: Kinematic inseparability is a state abstraction learnable via supervised learning that enables reward-free exploration.The section develops its definitions, properties, and intuition.
Backward kinematic inseparability: Backward kinematic inseparability is sufficient for exploration because equivalent observations induce the same ordering over policies by visitation probability.Thus, a policy maximizing visitation of an abstract state also maximizes visitation of each observation within it.
Forward kinematic inseparability: Forward kinematic inseparability requires identical transition probabilities to every next observation under every action.It captures forward dynamics that backward kinematic inseparability ignores and that are useful for model learning.
Kinematic inseparability: Kinematic inseparability combines backward and forward inseparability, and all three notions partition the observation space through equivalence relations.The associated dimensions NBD, NFD, and NKD are the sizes of the corresponding coarsest partitions.
Block MDP properties: For Block MDPs, observations emitted from the same state are kinematically inseparable, giving max{NFD, NBD} ≤ NKD ≤ |S|.The ideal case is NKD = |S|, though NKD may be smaller when distinct real states are indistinguishable under the abstraction.

4 Learning Kinematic Inseparability for Strategic Exploration

HOMER learns kinematic inseparability abstractions while conducting reward-free strategic exploration, producing policy covers and time-step-specific state abstractions. It uses contrastive supervised learning to guide exploration before optionally performing reward-sensitive optimization.

Algorithm overview: HOMER jointly learns a size-N policy cover and a state abstraction function for each time step during reward-free strategic exploration.The theoretical analysis assumes N ≥ N_KD and η ≤ η_min.
Optimization phases: After learning covers for all time steps, HOMER performs reward-sensitive optimization, or returns the covers directly for purely reward-free operation.This overall structure places a reward-free phase before the reward-sensitive phase.
Abstraction learning: At each time step, HOMER estimates an abstraction with contrastive classification using real and independently sampled imposter transitions.The dataset contains transitions generated by rolling in with a uniformly sampled prior policy, selecting a random action, and sampling the next observation.
Strategic exploration: HOMER uses the learned backward abstraction to define internal rewards, then invokes PSDP to construct the policy cover for the current time step.The learned forward abstraction is instead useful for auxiliary tasks such as transition-dynamics learning and visualization.

5 Theoretical Analysis

Theoretical guarantees show that ExpOracle and HOMER learn policy covers and reward-sensitive near-optimal policies sample-efficiently under realizability and policy-completeness assumptions. HOMER additionally learns the backward kinematic inseparability abstraction, with only logarithmic dependence on its potentially exponential hypothesis class.

Assumptions: The guarantees require realizability and policy completeness for the model classes Π and F.The realizability condition for F is aligned with using F to distinguish real from imposter transitions, requiring the Bayes-optimal classifier for the induced learning problems.
ExpOracle guarantee: ExpOracle outputs a policy cover and reward-sensitive policy with high probability, assuming access to a backward KI abstraction.Its sample complexity is polynomial in N, H, |A|, η_min^-1, ϵ^-1, and log |Π|/δ.
HOMER guarantee: HOMER achieves a similar guarantee without prior access to a backward KI abstraction by learning the abstraction and using it for exploration and policy optimization.The theorem applies to Block MDPs with η ≤ η_min and N ≥ N_KD.
HOMER guarantee: HOMER’s guarantee depends logarithmically on |Φ_N|, allowing an exponentially large abstraction class while retaining a guarantee qualitatively comparable to ExpOracle.This supports learning a backward KI abstraction from an exponentially large class and using it for exploration and policy optimization.
Computational efficiency: HOMER is computationally efficient in an oracle model when contextual bandit problems over Π and regression problems over F_N can be solved efficiently.The implementation’s most expensive component is the O(NH) calls to PSDP for learning policy covers.

6 Can We Use Existing State Abstraction Oracles?

Previously studied state abstractions generally do not support provably efficient exploration in general Block-MDPs when combined with a HOMER-like algorithm. The failures arise from representation collapse, policy-optimization tie-breaking, or reconstruction objectives, while the analysis does not rule out success in other settings or algorithms.

Summary: Overall, these approaches cannot ensure provably efficient exploration in general Block-MDPs within a HOMER-like algorithm, but may remain useful in other settings or algorithms.The conclusion is worst-case and does not imply that the approaches are ineffective universally.
Predicting Previous Action from Observations: Previous-action prediction can collapse distinct observations and fails to guarantee state-space coverage under a PSDP-style routine.A Bayes-optimal predictor can identify the previous action without distinguishing observations from states such as {s3, s4}, producing unfavorable policy-optimization tie-breaking.
Predicting Previous Action and Abstract State: Jointly predicting previous action and abstract state still collapses relevant state groups and fails in general Block-MDPs.Although this approach provably explores a restricted class of Block-MDPs, Figure 4a yields the same policy-optimization failure.
Predicting Previous Action and Abstract State: A model-based policy-cover method over abstract states avoids the Figure 4a tie-breaking issue but fails on Figure 4b.Because stochastic-start states cannot be separated, policies defined over abstract states must take the same action in s1 and s2, limiting reachability.
Training Autoencoders: Autoencoder reconstruction may fail to learn a meaningful representation even in a two-state problem with d-dimensional binary observations.The representation and decoder are trained by minimizing domain-specific reconstruction loss over raw observations.

7 Related Work

Prior work studies efficient exploration in small-state MDPs and rich-observation reinforcement learning, but existing approaches rely on restrictive assumptions, lack guarantees, or remain poorly understood. HOMER is most closely related to state-abstraction work while additionally addressing exploration and sample complexity.

Theoretical exploration: Small-state MDP methods achieve near-optimal sample complexity but do not exploit latent structure and therefore cannot scale to rich-observation environments.This line began with Kearns and Singh’s breakthrough result.
Theoretical exploration: Rich-observation theory establishes information-theoretic explorability and efficient algorithms in special settings, including deterministic or limited-stochastic latent-state dynamics.The present work continues this line while eliminating assumptions required by those settings.
Empirical exploration: Empirical exploration methods use pseudo-counts, optimism, intrinsic motivation, or prediction errors, but their behavior and failure modes remain poorly understood.The paper’s examples show that using their learned representations for provably efficient exploration is challenging.
State abstraction: Nachum et al. learn state abstractions with a similar supervised objective but do not address exploration or provide sample complexity guarantees.HOMER’s objective instead targets kinematic inseparability, motivated by the ExpOracle algorithm.

8 Proof of Concept Experiments

HOMER is evaluated on a high-dimensional, sparse-reward combination-lock task designed to defeat random exploration and standard reinforcement-learning baselines. Across experiments, HOMER learns effective policies, maintains visitation over time, and recovers the environment’s latent transition structure.

Environment: 10^-100: With H = 100 and K = 10, uniformly random actions receive the optimal reward with probability 10−100.The task combines high-dimensional observations, precarious dynamics, anti-shaped rewards, and extremely sparse optimal rewards.
Results: A2C and PPO fail for H > 3, while PCID fails for every tested horizon; HOMER learns a policy achieving half the optimal value.The baselines remain trapped by the anti-shaped reward or cannot separate observations sufficiently for exploration.
Results: None of the baselines learn the optimal policy at H = 100, whereas HOMER’s returns rise through three plateaus during training.Figure 6b reports average return per episode throughout training, with optimal policy value equal to 1.
Exploration: HOMER’s visitation probabilities do not decay significantly with time step, unlike those of the baselines.The comparison uses visitation counts from 100,000 uniformly sampled episodes.

9 Conclusion

HOMER is a model-free reinforcement learning algorithm for rich-observation environments that theoretically learns policy covers and near-optimal policies in Block MDPs with polynomial resources, while empirically solving difficult exploration and recovering latent transition dynamics.

Conclusion: HOMER learns both a policy cover and a near optimal policy for any Block MDP using polynomially many trajectories and polynomial time.These guarantees hold under certain computational and expressivity assumptions.
Conclusion: Empirically, HOMER solves a difficult exploration problem in a tractable manner and recovers the latent transition dynamics.

Appendices · A Properties of Homing Policies · B Properties of Kinematic Inseparability

The appendices establish foundational properties of homing policies and kinematic inseparability. They show that deterministic homing policies exist, composition can fail, and backward KI supports policy-cover construction through preserved visitation order and shared optimizers.

Appendices: The appendices organize notation and provide analyses of homing policies, kinematic inseparability, PSDP, ExpOracle, HOMER, supporting results, and experiments.The appendix overview explicitly lists these components and refers readers to Table 1 for notation.
A Properties of Homing Policies: For every state, a deterministic non-stationary policy exists that is a homing policy for that state.This result motivates restricting experimental searches to deterministic non-stationary policies.
A Properties of Homing Policies: Homing policies do not generally compose: appending a one-step policy to homing policies for earlier states may fail even with unrestricted policies.This non-compositionality justifies using the global PSDP policy-search procedure to find homing policies.
A Properties of Homing Policies: A constructed example shows that the homing policy for state s5 takes action a1 in s1 and reaches s5 with probability 1, unlike homing policies for s2, s3, and s4.The example demonstrates why earlier-state homing policies cannot simply be composed to obtain later-state homing policies.
B Properties of Kinematic Inseparability: Forward KI, backward KI, and KI each define an equivalence relation on observations.The proof follows from reflexivity, symmetry, and transitivity of the defining equality relations.
B Properties of Kinematic Inseparability: Observations sharing the same latent state are KI, and forward or backward KI observations must occur at the same time step.The Block MDP structure makes their transition distributions identical, while observations at different time steps have disjoint successor time layers.
B Properties of Kinematic Inseparability: Backward KI preserves ordering between policy visitation probabilities, enabling the backward KI relationship to identify a policy cover.For backward-KI observation sets, visitation-probability comparisons do not depend on the particular policy used in the comparison.
B Properties of Kinematic Inseparability: A single policy simultaneously maximizes visitation probability for every observation in a backward-KI set and optimizes the internal reward R′(x, a) := 1{x ∈X ′}.The common optimizer follows from the preserved visitation ordering and is optimal for rewarding membership in the observation set.

C Analysis of Policy Search by Dynamic Programming

This section analyzes PSDP’s statistical and computational guarantees. PSDP learns a sequence of policies through contextual bandit problems and runs in polynomial time with h contextual-bandit-oracle calls.

Theorem 5: PSDP returns policies under an α-policy-cover assumption, with the theorem guaranteeing its stated performance bound with probability at least 1 − hδ.The theorem applies when the policy covers have size at most N and under the specified internal- or external-reward conditions and Assumption 2.
Computational analysis: The algorithm runs in polynomial time with h calls to the contextual bandit oracle.
Contextual bandit guarantee: The contextual bandit oracle’s learned policy satisfies its performance guarantee with probability at least 1 − δ.The guarantee is established through a standard generalization bound for contextual bandit data collected under uniform logging.
Policy search: PSDP solves h contextual bandit problems sequentially to learn policies ˆπt for t = h, . . . , 1.Each problem is induced by a distribution generated by uniformly selecting a policy from the corresponding policy cover and executing it to visit the current state.
Proof of Theorem 5: The PSDP analysis combines the per-stage contextual-bandit guarantees with the performance difference lemma and the policy-cover property to prove the theorem.

D Analysis of the ExpOracle Algorithm … E.2 Building the policy cover

The analysis first shows that ExpOracle can construct policy covers and optimize rewards given a backward kinematic inseparability abstraction, then establishes that HOMER learns such abstractions through contrastive estimation and uses them to build policy covers inductively.

D Analysis of the ExpOracle Algorithm: ExpOracle uses a given backward kinematic inseparability abstraction to obtain policy covers and a reward-sensitive policy with high-probability guarantees.Its proof constructs covers inductively over time and then applies the supervised policy-learning theorem to the learned cover.
D Analysis of the ExpOracle Algorithm: For each abstraction index and time step, an internal reward identifies observations sharing that abstraction, so PSDP learns an approximate homing policy for them simultaneously.Repeating this over indices and time steps yields a cover for observations and latent states.
D Analysis of the ExpOracle Algorithm: When N ≥ NKD, realizability guarantees a function in the decoding class whose induced abstraction groups backward kinematically inseparable observations, allowing the ExpOracle analysis to apply.The resulting abstraction need not itself lie in ΦN if each time-slice component does and the proof’s two required properties remain valid.
E Analysis for the HOMER algorithm: HOMER constructs its policy cover inductively from h = 1 to h = H, using previously learned covers to produce the cover at the next time step.The base case sets Ψ1 = ∅ because the starting distribution itself provides the required coverage.
E.1 The supervised learning problem and a coupling: The supervised problem samples matched and independently generated transitions from policies in the previous cover, and an α-policy cover ensures ρh(s) ≥ αη(s) for every state.This lower bound provides coverage for learning the contrastive predictor.
E.1 The supervised learning problem and a coupling: The Bayes-optimal square-loss predictor is a transition-probability ratio involving the latent states and the cover-induced marginal, while empirical risk minimization supplies a finite-sample predictor in FN.The analysis couples the supervised distribution to show that the learned bottleneck abstraction approximately captures backward kinematic inseparability through its excess risk.
E.1 The supervised learning problem and a coupling: The learned abstraction is constant on kinematically inseparable observations in both the predecessor and successor arguments of the Bayes predictor.This establishes realizability of the bottleneck structure for the supervised learning problem.
E.2 Building the policy cover: Using internal rewards induced by the learned abstraction, PSDP produces a policy cover for the current latent-state layer under a comparison inequality between visitation probabilities.The cover guarantee holds when earlier layers have α-policy covers and the N PSDP calls and regression guarantee succeed.

E.3 Wrapping up the proof

Lemma 12 supplies the core coverage guarantee, and completing the induction shows that the learned policies form 1/2-policy covers when η ≤ ηmin. These covers then enable the reward-sensitive PAC guarantee, with sample and computational complexity determined by the stated trajectory and oracle-call costs.

The policy covers: Lemma 12 is the core technical result certifying good coverage at each time h, after which only induction, parameter setting, and a union bound remain.The proof invokes Corollary 9 and Theorem 5 across the horizon and exploration calls before accounting for the total failure probability.
The policy covers: When η ≤ ηmin, the chosen npsdp and nreg values certify that Ψh is a 1/2-policy cover for Sh, completing the inductive policy-cover guarantee.The algorithm sets these quantities using η rather than ηmin, which need not be known.
The reward sensitive step: A single PSDP call with the external reward R, together with Theorem 5 and the selected neval, yields the PAC guarantee.The failure probability for this reward-sensitive step has already been included in the proof accounting.
Sample complexity: The trajectory sample complexity combines H supervised learning problems, NH PSDP calls using npsdp, and one PSDP call using neval.The resulting sample complexity is measured in trajectories.
Computational complexity: The running time is HT + HTimereg(nreg) + NH2Timepol(npsdp) + HTimepol(neval), combining trajectory collection with oracle-call time.Here T denotes the number of trajectories.

F Supporting results

The supporting results establish a performance difference identity, a Bernstein concentration tool, and a square-loss learning guarantee under realizability and parametric metric entropy. The square-loss bound is presented as a standard excess-risk result for well-specified infinite function classes.

Performance difference lemma: The performance difference lemma relates the values of any two non-stationary policies for any episodic decision process and reward function.Its proof uses a standard telescoping argument.
Bernstein’s inequality: Bernstein’s inequality provides a high-probability bound for independent, bounded, zero-mean random variables using their total variance.The proposition assumes |U_t| ≤ R almost surely and bounds the result with probability at least 1 − δ.
Square-loss learning: Under realizability and parametric pointwise metric entropy growth, empirical risk minimization over G satisfies a high-probability square-loss excess-risk guarantee.The result applies to samples from a distribution over examples and labels in [0, 1], with a universal constant C > 0.
Square-loss learning: The square-loss result is described as a standard bound for well-specified infinite function classes, while sharper complexity-based guarantees can be substituted compatibly.The alternatives include empirical covering numbers, combinatorial parameters, and localization.

G Experimental Setup and Optimization Details

The experiments implement HOMER’s oracle and policy-learning subroutines with supervised optimization, while comparing against PCID and several neural reinforcement-learning baselines under specified training protocols. HOMER also uses empirical computational and statistical optimizations, with implementation details and hyperparameters reported for each method.

HOMER: HOMER implements REG as supervised binary classification, trains with Adam mini-batches, and selects the best model using validation performance.Models are trained for a maximum number of epochs, with bottleneck models initialized from models trained without the bottleneck.
HOMER: HOMER learns CB policies by predicting immediate rewards with mean squared error, equivalent to one-step Q-learning, and acting greedily with arg max_a Qθ(x, a).CB uses Adam mini-batches, no validation dataset, and the model from the final epoch.
HOMER: HOMER uses two empirical optimizations described in Section 5, which provide computational and statistical advantages.The diabolical combination lock hyperparameters are reported in Table 2, and the implementation uses PyTorch 1.1.0.
PCID: PCID uses publicly released code, predicts the previous state and action, clusters predicted probabilities with k-means, and uses a linear model sufficient for the combination-lock task.The evaluation follows the authors’ recommended hyperparameters, with one exception concerning data collection.
Baseline training: A2C, PPO, A2C + RND, and PPO + RND baselines train for a maximum of 10 million episodes using 2-layer, 64-unit ReLU MLPs and RMSProp.Learning rates are tuned over {0.01, 0.001, 0.0001}; RND networks use 2-layer, 100-unit ReLU MLPs.