Factorized Spectral Representations for Reinforcement Learning

2026-07-15

AI Summary

This work addresses the problem of efficiently learning compact world model representations for reinforcement learning from interaction data. The proposed method models the environment transition kernel as a third-order tensor over state–action–next-state triples and introduces, for the first time, CP decomposition to obtain factorized feature maps for each modality. A unified spectral representation is constructed by jointly optimizing encoders via a noise-contrastive objective. This approach substantially reduces the hypothesis space, leading to improved sample efficiency and strong performance on high-dimensional control tasks. Moreover, the learned state encoder exhibits cross-actuator transferability, requiring only fine-tuning of the action encoder to adapt to new dynamics.

Abstract

Learning a compact model of the world from interaction data is central to sample-efficient deep reinforcement learning. Spectral representation methods have become the leading paradigm for representation learning in continuous control by taking a matrix view of the transition kernel, with state-action pairs on one side and next states on the other, and learning a low-rank factorization through self-supervised contrastive objectives. We take this view one step further. The transition kernel is naturally a three-mode tensor over states, actions, and next states, and a CP decomposition gives one feature map per mode. We propose FaStR, which fits this decomposition with a noise contrastive objective, producing separate state, action, and next-state encoders that together form a single spectral representation. The factored form yields a smaller hypothesis class, and the sample size needed for representation learning shrinks by a factor that scales with the smaller of the state and action dimensions. Empirically, FaStR delivers its largest gains on high-dimensional locomotion tasks whose dynamics align with the factored structure, and the learned state encoder transfers intact across actuator shift while only the action encoder is retrained.

Problem

• spectral representation
• reinforcement learning
• transition kernel
• sample efficiency
• representation learning

Innovation

• tensor decomposition
• spectral representation
• reinforcement learning
• CP decomposition
• sample efficiency