PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Good State and Action Representations via Tensor Decomposition

Chengzhuo Ni, Yaqi Duan, Munther Dahleh, Anru Zhang, Mengdi Wang

arXiv:2105.01136v2stat.MLcs.LG

TL;DR

High-dimensional state and action spaces make useful RL representations challenging: naive grid discretization can cause exponentially many states, and states and actions are entangled in trajectories. The paper learns low-dimensional embeddings from empirical transitions using kernelized tensor approximation; its guarantees include at-most-quadratic growth of H-step state abstraction error and recovery of latent blocks in a simulation. The paper identifies prior knowledge of tensor rank as a theoretical assumption.

  • Problem

    In high-dimensional RL, grid discretization can produce exponentially many states, while state and action structure is entangled in trajectories despite potentially separate low-dimensional structure.

    Source: §1 · pp. 1–2

  • Method

    The method uses kernel function approximation, importance sampling, mean embedding, and low-rank tensor approximation to learn state and action representations from empirical trajectories.

    Source: §1.1 · p. 3

  • Results

    The H-step state abstraction error grows at most quadratically with horizon, and simulation clustering validates recovery of latent state and action blocks.

    Source: §4.3 · p. 12, §5 · p. 14

  • Takeaways & Limitations

    The authors conclude that learned state and action abstractions are useful for approximate policy evaluation.

    Source: §4.3 · p. 12

    The theory assumes prior knowledge of tensor rank, which is hard to estimate without additional strong assumptions such as a bounded eigengap.

    Source: §3.2 · p. 8

Abstract

from arXiv · show

The transition kernel of a continuous-state-action Markov decision process (MDP) admits a natural tensor structure. This paper proposes a tensor-inspired unsupervised learning method to identify meaningful low-dimensional state and action representations from empirical trajectories. The method exploits the MDP's tensor structure by kernelization, importance sampling and low-Tucker-rank approximation. This method can be further used to cluster states and actions respectively and find the best discrete MDP abstraction. We provide sharp statistical error bounds for tensor concentration and the preservation of diffusion distance after embedding. We further prove that the learned state/action abstractions provide accurate approximations to latent block structures if they exist, enabling function approximation in downstream tasks such as policy evaluation.

1. Introduction

High-dimensional RL needs useful abstractions because grid discretization can create exponentially many states, while states and actions have distinct structures entangled in trajectories. The paper proposes learning low-dimensional state and action representations from empirical trajectories using the MDP’s tensor structure.

1. Introduction

  • Grid discretization can create exponentially many states, and useful representations support function approximation that generalizes from seen to unseen states.Large action spaces further complicate learning because states and actions with distinct low-dimensional structures are entangled in trajectories.Source: §1 · pp. 1–2

1.4 Tensor and Tucker Decomposition

  • Tucker decomposition expresses a tensor through a smaller core and mode-specific matrices, while low Tucker rank connects transition structure to state and action abstractions.The paper relates this structure to block MDPs and soft state aggregation.Source: §1.1 · p. 2, §1.4 · p. 5

1.1 Our Approach

  • The method learns state and action embeddings from empirical transitions using kernelized diffusion distance, importance sampling, mean embedding, and low-rank tensor approximation.The paper also proposes using the embeddings to find discrete MDP abstractions and recover latent block structure.Source: §1.1 · pp. 2–3
  • The embedding maps are intended to preserve state and action predictability in kernelized diffusion distance, with guarantees established through tensor concentration analysis.The paper presents this as a theoretical guarantee for the learned representations.Source: §1.1 · p. 3
  • The proposed discrete state/action abstraction method is claimed to recover latent block structures of aggregable MDPs.The paper also claims guarantees for approximating transition and value functions, though the passage does not give the full error statement.Source: §1.1 · p. 3
  • Simulations are reported to show an advantage over vanilla and top-r kernel PCA baselines.The passage does not specify metrics or numerical comparisons.Source: §1.1 · p. 3

1.2 Related Literature

  • Prior RL representation methods often require known problem structure or a covering finite feature class; earlier tensor methods cited here target finite state-action MDPs with low CP rank.Source: §1.2 · p. 4

1.3 Markov Decision Process

  • An MDP is specified by state and action spaces, a transition kernel, and a one-step reward function; a policy induces a Markov chain.Source: §1.3 · p. 4

2. A Tensor View of Markov Decision Process

The paper represents a continuous-state MDP’s transition operator with a low-Tucker-rank structure in state, action, and next-state dimensions. Kernel function approximation turns this structure into a finite-dimensional transition tensor, which relates to several reduced-order MDP models.

2. A Tensor View of Markov Decision Process

  • The transition kernel is modeled in state and action RKHSs, with a low-Tucker-rank assumption capturing compressible state/action structure that preserves dynamics.The paper relates this assumption to soft state aggregation, rich-observation MDPs, contextual MDPs, linear/factor MDPs, and kernel MDPs.Source: §2 · pp. 5–6
  • Known kernel functions are approximated with finitely many features, allowing the continuous transition problem to be represented by a finite-dimensional tensor.The paper says it applies to arbitrary state and action spaces when suitable kernels are available.Source: §2 · p. 6
  • The paper states that low Tucker rank also holds for block MDPs and soft aggregation, and that transition-model rank is independent of the reward function.Source: §2 · p. 6

3. Tensor-Inspired State and Action Embedding Learning

This section develops tensor-based, decoupled state and action embeddings from empirical transitions, using importance sampling and low-rank estimation. The embeddings target predictive similarity and are supported by bounds on diffusion-distance preservation and sample requirements.

3.1 Tensor MDP Mean Embedding by Importance Sampling

  • Importance sampling decouples states from actions; without it, behavior-policy coupling can raise tensor ranks and prevent recovery of the best low-dimensional structure.The data consist of transitions from a single sample path generated by a known behavior policy.Source: §3.1 · pp. 6–7

3.2 Low-Rank Estimation of Transition Tensor

  • The method estimates a Tucker-low-rank approximation of the mean embedding tensor, then uses it to estimate the transition operator.The theory assumes prior knowledge of tensor rank; practice can tune rank using scree-plot elbows and cross-validation.Source: §3.2 · pp. 7–8

3.3 Learning State and Action Embeddings

  • State and action embedding maps are obtained from factors of the estimated transition tensor’s Tucker decomposition.The maps define joint embeddings used to represent predictive similarity through diffusion distance.Source: §3.3 · pp. 8–9
  • Under Assumptions 1–2 and orthogonal-basis conditions, embedded and original diffusion distances differ by at most 2ϵ.The bound holds for all state-action pairs.Source: §3.3 · p. 9
  • ˜Ω(dAdSr) samples are required for joint treatment, versus ˜Ω((dX + dA)r) for the tensor approach when Tucker ranks are equal.The comparison assumes state and action feature dimensions dS and dA before embedding, with r = l = m.Source: §3.3 · p. 9

4. Estimating the Optimal Discrete MDP Abstraction

The section formulates discrete state-action abstraction as a partitioning problem and approximates it by separately clustering learned state and action representations. Under the stated assumptions, the analysis bounds clustering and policy-evaluation errors for the resulting abstractions.

4.1 Optimal Partition of State and Action Spaces

  • The optimal abstraction partitions state and action spaces into blocks whose transition distributions approximate one-step expected transitions for functions in HS.The objective is to choose the blocks and corresponding transition distributions to minimize this approximation error.Source: §4.1 · p. 10

4.2 Decoupled State and Action Clustering

  • The empirical clustering problem uses kernelized diffusion distance to approximate the optimal partition when the transition density is unknown.A relaxation yields two simpler subproblems, and k-means-type algorithms compute decoupled state and action clusters from the learned representations.Source: §4.2 · pp. 10–11

4.3 Theoretical Guarantee

  • With Assumptions 1–3, the paper gives a high-probability misclassification-error bound for recovering latent state and action blocks when the sample size is sufficiently large.The bound includes a factor proportional to the product rlm of the transition tensor’s Tucker ranks, which can be large even when each rank is individually small.Source: §4.3 · p. 12
  • The H-step policy-evaluation error from inaccurate state abstraction grows at most quadratically with H, rather than exponentially.The authors conclude that this error does not suffer from the curse of horizon, supporting approximate policy evaluation with the learned abstractions.Source: §4.3 · p. 12

5. Numerical Experiment

In a controlled stochastic-process MDP, the method recovers the planted state and action blocks. In transition-tensor estimation, it generally outperforms the vanilla and top-r baselines, with exceptions at very small sample sizes and high rank constraints.

State-action Clustering

  • The state and action spaces are both R2; after pair (sk, ak), the next state is Xτ(k+1) for some τ > 0, and the behavior policy always chooses actions from a standard normal distribution.Source: §5 · p. 13
  • The learned state and action clusters reveal the latent blocks in the controlled stochastic-process MDP.The experiment uses a two-dimensional state and action space, with a wavy potential and block-wise constant control function.Source: §5 · pp. 13–14, Figure 2 · p. 13

Low-Rank Estimation of the Transition Tensor

  • The low-tensor-rank method consistently outperforms the vanilla and top-r baselines for most settings, while top-r performs slightly better when n is very small.The three methods have similar performance at rank constraint (60, 30, 60) or higher; the authors attribute this to the constraint approaching the feature dimensions and added computational noise.Source: Figure 3 · p. 14, §5 · p. 14

A. The HOOI Algorithm

The appendix presents HOOI as the algorithm for low-rank MDP tensor decomposition, using the tensor mean embedding and Tucker-rank choices to construct an estimated tensor.

A. The HOOI Algorithm

  • Algorithm 3 applies HOOI to the tensor mean embedding using Tucker-rank choices (r, l, m) and an iteration limit tmax.Its stated output is an estimated tensor formed using the iterated factor matrices.Source: A. The HOOI Algorithm · p. 20
  • The algorithm uses SVDr to return the leading r singular vectors of a matrix.This operation is specified as part of the algorithm’s initialization.Source: A. The HOOI Algorithm · p. 20

B. Examples of Low Rank MDPs

The examples connect low Tucker rank to block and latent-state-action MDP structures. They also illustrate why preserving the state-action-state tensor can retain aggregation information that state-transition matrix methods lose.

B. Examples of Low Rank MDPs

  • In a block MDP, state and action abstractions imply transition-tensor Tucker rank at most (|S̃|, |Ã|, |S|).This is the hard-aggregation example.Source: B. Examples of Low Rank MDPs · p. 20, Figure 4 · p. 20
  • A latent-state-action MDP with latent dimensions r, l, and m has transition-kernel Tucker rank (r, l, m).The example describes soft aggregation through latent state, action, and next-state variables.Source: B. Examples of Low Rank MDPs · p. 21
  • For the illustrated MDP, state-transition matrix methods average away the meta-state partition, whereas the tensor formulation retains it and allows reliable state aggregation.The example uses meta-states {1, 2} and {3, 4} under a random policy.Source: B. Examples of Low Rank MDPs · p. 21

C. Derivation of Optimization Problem (4)

This section derives optimization problem (4) by first choosing the best feature-related term for fixed matrices, then formalizing the resulting problem and approximating it with empirical data.

C. Derivation of Optimization Problem (4)

D. Experiment Details

The experiment uses Gaussian kernels and random Fourier features, then evaluates learned state and action abstractions; the reported comparison says top-r clustering does not capture the transition kernel's correct clustering information.

D. Experiment Details

  • The experiment uses Gaussian kernels with τ = 0.1 and σ = 0.5.The state and action kernels use random Fourier features; action features are orthogonalized with respect to L2(η).Source: D. Experiment Details · pp. 22–23
  • The top-r method does not capture the transition kernel's correct clustering information compared with the authors' method.Figure 5 shows learned state and action abstractions across clustering sizes for top-r.Source: D. Experiment Details · p. 23, Figure 5 · p. 23
  • For clustering, the experiment uses n = 10^6 and (r, l, m) = (3, 3, 3), with state features orthogonalized with respect to L2(ξ).The state and action features are also generated using Ns = 100 and Na = 50 random Fourier features, respectively.Source: D. Experiment Details · p. 23
  • For estimation, the ground truth is approximately obtained using the vanilla method with sample size n = 10^6.Source: D. Experiment Details · p. 23

E. Technical Lemmas

The technical lemmas address Tucker-rank tensor representations, mixing-time conditions, tensor concentration, and comparison of Tucker decompositions; they also invoke sphere nets.

E. Technical Lemmas

  • A tensor with Tucker rank bounded by (r_1, r_2, …, r_N) has column-wise orthonormal factor matrices.Source: E. Technical Lemmas · p. 23
  • One lemma states a condition involving worst-case MDP mixing time, a policy's invariant distribution, and an initial distribution.The supplied statement does not include the full condition or bound.Source: E. Technical Lemmas · p. 24
  • The appendix gives an ε-net size bound of (1 + 2/ε)^p for the sphere S^(p−1).Another lemma relates three sphere nets to a tensor spectral-norm quantity, but its bound is not shown here.Source: E. Technical Lemmas · p. 24
  • A separate lemma addresses concentration in tensor spectral norm under Assumptions 1–2.The passage does not supply the concentration bound.Source: E. Technical Lemmas · p. 24
  • A lemma compares two same-dimension order-3 tensors represented by Tucker decompositions with core tensors and factor matrices.The supplied passage does not state the resulting comparison.Source: E. Technical Lemmas · p. 25

F. Proofs

The proofs establish low-Tucker-rank transition representations and derive concentration, perturbation, and representation-error guarantees. They bound distortion in embedded diffusion distances, show recovery of state-action blocks up to relabeling, and connect block and transition discrepancies to downstream value approximation.

Proofs of Lemmas 1–2

  • The transition representation has Tucker rank at most (r, l, m), the dimensions of its latent core.The proof rewrites the conditional next-state feature expectation as a Tucker-form tensor contraction.Source: F. Proofs · p. 26

Proofs of Lemmas 15–17

  • Mixing-time blocking, matrix Bernstein bounds, and net-based union bounds establish concentration for covariance and transition-tensor estimates.The proof then relates covariance estimation error to its inverse and bounds the final tensor-estimation error.Source: F. Proofs · pp. 31, 33, 34, 35, 37, 38

Proof of Lemma 18

  • Wedin’s lemma bounds the estimated third-mode subspace error by the transition-tensor estimation error.The proof applies the perturbation bound uniformly over unit vectors in the first mode.Source: F. Proofs · p. 40

Proof of Theorem 5

  • Theorem 5 bounds diffusion-distance distortion by the feature-embedding errors for the two state-action pairs.The proof controls those errors using the preceding subspace perturbation result.Source: F. Proofs · pp. 40–41

Representation-error proof

  • Theorem 4 and Lemma 17 yield a high-probability representation-error bound, with the learned third-mode subspace aligned to the true one up to orthogonal rotation.The proof invokes the subspace-angle result to choose the rotation before concluding the bound.Source: F. Proofs · pp. 45–46

Block recovery proof

  • For sufficiently large n, the embedding separates points outside different blocks while remaining constant within each block, so recovered state and action labels match up to permutations.The proof rearranges the learned cluster indices to align them with the true indices.Source: F. Proofs · p. 48

Proofs of Theorem 8

  • Theorem 8’s proof uses Theorem 7’s label matching to bound the measure discrepancy between corresponding true and learned state-action blocks.The proof first aligns the block labels by permutations; its displayed bound depends on the mass outside the matched blocks.Source: F. Proofs · p. 49

Value-approximation proof

  • The value-approximation proof couples trajectories through an auxiliary mismatch indicator and bounds mismatch probability by H times the largest block-mass discrepancy.It embeds overlapping true and learned blocks in a discrete MDP with an absorbing mismatch state; rewards are block-dependent and bounded in [0, 1].Source: F. Proofs · pp. 50–52
  • On the overlapping blocks, the transition-distribution difference is bounded by the largest transition-entry error plus the difference between the compared state-action distributions.This decomposition supports the value comparison alongside the separate bound for probability mass outside matched blocks.Source: F. Proofs · p. 52
Loading 2105.01136v2…