A Tensor Network Implementation of Multi Agent Reinforcement Learning

TL;DR

The paper presents a tensor-network (TN) formulation for representing the expected return in finite multi-agent MDPs (FMDPs), addressing the curse of dimensionality that plagues MARL. By building a TN that encodes state transitions, rewards, actions, and policies—using mechanisms such as Matrix Product Operators for rewards, copy tensors for state propagation, and DMRG-style optimisation—the approach enables exact or approximate planning and policy learning in MARL. A two-agent random-walker demonstration validates the framework, showing correct policy optimisation and substantial tensor-size reductions (up to 97.5% lossless) via exact decompositions like SVD. The work highlights planning and model-learning as practical extensions and discusses scalability to more agents, as well as potential performance comparisons against traditional distribution-based MARL methods.

Abstract

Recently it has been shown that tensor networks (TNs) have the ability to represent the expected return of a single-agent finite Markov decision process (FMDP). The TN represents a distribution model, where all possible trajectories are considered. When extending these ideas to a multi-agent setting, distribution models suffer from the curse of dimensionality: the exponential relation between the number of possible trajectories and the number of agents. The key advantage of using TNs in this setting is that there exists a large number of established optimisation and decomposition techniques that are specific to TNs, that one can apply to ensure the most efficient representation is found. In this report, these methods are used to form a TN that represents the expected return of a multi-agent reinforcement learning (MARL) task. This model is then applied to a 2 agent random walker example, where it was shown that the policy is correctly optimised using a DMRG technique. Finally, I demonstrate the use of an exact decomposition technique, reducing the number of elements in the tensors by 97.5%, without experiencing any loss of information.

Figures (16)

• Figure 1: The classical graphical representations of different ranked tensors. The number of lines coming off the green shape indicates how many indices the tensor has (and therefore its rank). In this report, these shapes will only be followed loosely.
• Figure 2: A figure showing the contraction of a rank 3 tensor with a rank 4 tensor. The indices are labelled to make it clear what the shared indices are, which in this case are $i_{2}$ and $i_{3}$.
• Figure 3: Showing how the tensor network is contracted leaving $\bm{\pi}_{3}$ being contracted with another object termed $\mathbf{Z}_{3}$. Note the order of contraction, which involves contracting the first two layers, and the last layer individually. These layers are then contracted with the third layer. This order is superior to just contracting in time order, as it ensures that the ranks of resulting tensors are minimised at any time throughout the contractions.
• Figure 4: A figure displaying how the normal distribution is discretised. The blue area between -1 and 1 represents the probability of the discrete variable $P(\mathcal{X}_{d} = 0)$. Due to the symmetry of the normal distribution, the area of the red and green regions are equal, representing the fact that $P(\mathcal{X}_{d} = -1) = P(\mathcal{X}_{d} = 1)$. If $\sigma = 1$, $P(\mathcal{X}_{d} = \pm1) \approx 0.16$ and $P(\mathcal{X}_{d} = 0) \approx 0.68$.
• Figure 5: Showing 100 trajectories that were generated before and after policy optimisation. Red trajectories represent those that did not achieve the maximum return, and blue represents those that did. As there are significantly more red trajectories, they are plotted slightly transparently. Also shown is the expected return. a) There is a large negative expected return, giving rise to a large number of trajectories that do not obey the objective. b) The expected return is now 1, which is the highest possible score achievable in this simulation, and every trajectory obeys the objective.