Multilinear Tensor Low-Rank Approximation for Policy-Gradient Methods in Reinforcement Learning

TL;DR

This work introduces multilinear tensor low-rank policy models for policy-gradient reinforcement learning by organizing policy parameters into a PARAFAC-decomposed tensor, greatly reducing parameter counts from a full tensor to a sum of mode sizes times the rank. The framework yields tensor policy models for Gaussian and softmax policies and integrates them into PG, actor-critic, TRPO/PPO variants, and their trust-region counterparts, with convergence guarantees under mild assumptions. Empirical results across classical control tasks and a wireless-communications setup show that the tensor-based methods achieve comparable rewards with faster convergence and far fewer parameters than neural-network baselines. Overall, the approach leverages intrinsic low-rank structure in RL representations to improve efficiency and scalability without sacrificing performance.

Abstract

Reinforcement learning (RL) aims to estimate the action to take given a (time-varying) state, with the goal of maximizing a cumulative reward function. Predominantly, there are two families of algorithms to solve RL problems: value-based and policy-based methods, with the latter designed to learn a probabilistic parametric policy from states to actions. Most contemporary approaches implement this policy using a neural network (NN). However, NNs usually face issues related to convergence, architectural suitability, hyper-parameter selection, and underutilization of the redundancies of the state-action representations (e.g. locally similar states). This paper postulates multi-linear mappings to efficiently estimate the parameters of the RL policy. More precisely, we leverage the PARAFAC decomposition to design tensor low-rank policies. The key idea involves collecting the policy parameters into a tensor and leveraging tensor-completion techniques to enforce low rank. We establish theoretical guarantees of the proposed methods for various policy classes and validate their efficacy through numerical experiments. Specifically, we demonstrate that tensor low-rank policy models reduce computational and sample complexities in comparison to NN models while achieving similar rewards.

Figures (4)

• Figure 1: $\mathrm{NFE}$ between the tensor of Gaussian and softmax parameters obtained via NN-based PPO and their low-rank PARAFAC decomposition in the continuous and discrete variants of the Pendulum environment. The error decreases as the rank of the approximation increases.
• Figure 2: Median return per episode of tensor low-rank policies (TLR) against NN-based policies (NN) across $100$ experiments in various continuous action-space RL problems and algorithmic setups. The displayed confidence interval are the interquartile ranges. TLR policies converge faster than NN-based policies in almost all the scenarios. Furthermore, they require less parameters.
• Figure 3: Median return per episode of tensor low-rank policies (TLR) against NN-based policies (NN) across $100$ experiments in various discrete action-space RL problems and algorithmic setups. The displayed confidence interval are the interquartile ranges. TLR policies converge faster than NN-based policies in almost all the scenarios. Furthermore, they require less parameters.
• Figure 4: Median return per episode of tensor low-rank policies (TLR) against NN-based policies (NN) across $100$ experiments in the wireless communications setup. The displayed confidence interval are the interquartile ranges. PTLRPO converges faster than NN-PPO while requiring less parameters.

Theorems & Definitions (2)

• Theorem 1
• proof