[Paper Review] Learning Good State and Action Representations via Tensor Decomposition

Chengzhuo Ni, Yaqi Duan+3 more|arXiv (Cornell University)|May 3, 2021

Tensor decomposition and applications50 references4 citations

TL;DR

This paper proposes a tensor decomposition-based unsupervised method to learn low-dimensional state and action representations from empirical trajectories in continuous-state-action Markov decision processes (MDPs). By leveraging kernelization, importance sampling, and low-Tucker-rank approximation, the method identifies meaningful abstractions that preserve diffusion distances and accurately approximate latent block structures, enabling effective function approximation in downstream RL tasks with rigorous statistical error bounds.

ABSTRACT

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.

Motivation & Objective

• To address the challenge of learning meaningful low-dimensional representations of continuous states and actions in MDPs without supervision.
• To exploit the inherent tensor structure of MDP transition kernels for improved representation learning.
• To enable accurate abstraction of states and actions via clustering and discrete MDP construction.
• To provide statistical error bounds for tensor concentration and diffusion distance preservation after embedding.
• To ensure that learned representations approximate underlying latent block structures, supporting downstream function approximation.

Proposed method

• The method applies kernelization to map states and actions into a reproducing kernel Hilbert space, enabling tensor-based modeling of the transition kernel.
• Importance sampling is used to efficiently estimate the transition kernel from empirical trajectories, reducing variance in estimation.
• Low-Tucker-rank approximation is applied to the kernelized transition tensor to extract low-dimensional state and action representations.
• The algorithm clusters states and actions based on the learned representations to form discrete MDP abstractions.
• Statistical guarantees are derived for tensor concentration and preservation of diffusion distances in the embedded space.
• The approach is designed to recover latent block structures in the transition kernel when they exist, supporting function approximation.

Experimental results

Research questions

• RQ1 Can tensor decomposition techniques be effectively adapted to learn low-dimensional representations of continuous states and actions in MDPs?
• RQ2 How well do the learned representations preserve the geometric structure of the MDP, particularly diffusion distances?
• RQ3 To what extent can the method recover latent block structures in the transition kernel, enabling accurate function approximation?
• RQ4 What are the statistical error bounds for tensor concentration and representation fidelity in the proposed framework?
• RQ5 Can the method produce high-quality discrete MDP abstractions through clustering of the learned representations?

Key findings

• The method achieves sharp statistical error bounds for tensor concentration, ensuring reliable estimation of the kernelized transition tensor.
• The learned representations preserve diffusion distances with high accuracy, maintaining the intrinsic geometry of the MDP.
• When latent block structures exist in the transition kernel, the method recovers them with high fidelity, enabling effective function approximation.
• The low-Tucker-rank approximation enables compact, meaningful representations that support clustering of states and actions.
• Theoretical guarantees show that the method provides consistent and accurate abstractions even with limited trajectory data.
• Empirical results demonstrate that the learned representations lead to improved performance in downstream policy evaluation tasks.