# Model based Multi-agent Reinforcement Learning with Tensor Decompositions

Pascal Van Der Vaart
TU Delft
<EMAIL_ADDRESS>&Anuj Mahajan
University of Oxford
<EMAIL_ADDRESS>&Shimon Whiteson
University of Oxford
<EMAIL_ADDRESS>
###### Abstract

A challenge in multi-agent reinforcement learning is to be able to generalize
over intractable state-action spaces. Inspired from Tesseract [Mahajan et al.,
2021], this position paper investigates generalisation in state-action space
over unexplored state-action pairs by modelling the transition and reward
functions as tensors of low CP-rank. Initial experiments on synthetic MDPs
show that using tensor decompositions in a model-based reinforcement learning
algorithm can lead to much faster convergence if the true transition and
reward functions are indeed of low rank.

##  1 Introduction

Recent progress in multi-agent RL has been promising towards creating agents
which are capable of generalising over multiple tasks [DeepMind-OEL et al.,
2021], they have also demonstrated effectiveness in dealing with the problem
of exploration in a large action space [Mahajan et al., 2019, Gupta et al.,
2020] and overcoming intractability arising from its exponential growth in the
number of agents[Mahajan et al., 2021, Wang et al., 2020a, b] when learning
under constraints like decentralisation. Inspired from Tesseract [Mahajan et
al., 2021], which utilises tensor decomposition structure in factored action
spaces, we investigate whether tensor decompositions can be used to attain
generalisation across the state-action space in cooperative multi-agent
setting towards obtaining better sample efficiency. In this position paper, we
focus on the model based setting. Initial empirical results on randomly
generated MDPs provide promising evidence for state-action generalisation and
sample efficiency using tensor decompositions over baseline model based
algorithms which do not use the tensor approximation.

In multi-agent reinforcement learning, the goal is to find a policy for
multiple agents that performs well on a given task. Tasks are formalized by
Markov decision processes (MDP), which are described by a transition function
and a reward function. In the RL setting, the transition and reward functions
are unknown, and finding a policy that achieves high reward requires
exploration to gather data from the MDP and learn its dynamics. Because the
size of the action space grows exponentially with the amount of agents, it is
especially important in multi-agent reinforcement learning to learn with high
sample efficiency as coverage over all state action pairs is not feasible.
This work showcases how the use of low rank CP-decompositions can drastically
improve sample efficiency of classic model-based reinforcement learning
algorithms and provide better generalisation when the transition and reward
tensors are of low CP-rank.

##  2 Background

###  2.1 Tensor decompositions

###### Definition 1

An order $n$ tensor over a field $\mathbb{F}$ with dimensions
$d_{1},d_{2},\dots,d_{n}$ is a multilinear map
$T:\mathbb{F}^{d_{1}}\times\mathbb{F}^{d_{2}}\times\cdots\times\mathbb{F}^{d_{n}}\rightarrow\mathbb{F},$
and can be represented by a n-dimensional $d_{1}\times d_{2}\times\cdots\times
d_{n}$ array $T_{i_{1}i_{2}\dots i_{n}}$ such that the mapping is defined as

$T(u^{1},\dots,u^{n})=\sum_{i_{1}=1}^{d_{1}}\sum_{i_{2}=1}^{d_{2}}\dots\sum_{i_{n}=1}^{d_{n}}T_{i_{1}i_{2}\dots
i_{n}}u^{1}_{i_{1}}u^{2}_{i_{2}}\dots u^{n}_{i_{n}}$

for $u^{i}\in\mathbb{F}^{d_{i}}$. For a set of indices $i_{1},\dots
i_{k}\in\\{1,\dots,n\\}$ the expression

$T(u^{1},\dots,u^{i_{1}-1},I,u^{i_{1}+1},\dots
u^{i_{2}-1},I,u^{i_{2}+1},\dots,u^{i_{k}-1},I,u^{i_{k}+1},\dots,u^{n})$

denotes a $d_{i_{1}}\times\dots\times d_{i_{k}}$ tensor defined by the mapping
$(u^{i_{1}},u^{i_{2}},\dots,u^{i_{k}})\mapsto T(u^{1},\dots,u^{n})$.

The set of order n tensors with dimensions $d_{1},d_{2},\dots d_{n}$ over
$\mathbb{F}$ is denoted by $\mathbb{F}^{d_{1}\times d_{2}\times\cdots\times
d_{n}}$.

The CANDECOMP/PARAFAC (CP) decomposition for tensors can be thought of as a
generalization of the singular value decomposition for matrices. Related work
involving tensors and tensor decompositions can be found in appendix A.

###### Definition 2

A rank $r$ CP-decomposition of a tensor $T\in\mathbb{F}^{d_{1}\times
d_{2}\times\cdots\times d_{n}}$ is a set of vectors $(u_{l}^{i})^{i=1,\dots
n}_{l=1,\dots r},u_{i}^{l}\in\mathbb{F}^{d_{i}}$ and scalars
$(w_{l})_{l=1,\dots,r}\in\mathbb{F}$ such that

$T=\sum_{l=1}^{r}w_{l}u_{l}^{1}\otimes u_{l}^{2}\otimes\cdots\otimes
u_{l}^{n},$

where $\|u_{l}^{i}\|=1$ for all $l\in\\{1,\dots,r\\}$ and
$i\in\\{1,\dots,n\\}$.

The tensor $T$ is said to be of CP-rank $r$ if $r$ is the smallest number for
which a rank $r$ CP-decomposition for $T$ exists.

It is clear that if a large $n\times n\times n$ tensor $T$ has low rank, the
search space of an application which requires an estimate of $T$ can be
greatly reduced by incorporating the low rank information. Instead of
estimating $n^{3}$ parameters, the problem can be described by $3rn$
parameters in decomposed form instead. The main idea of this work is to use
this fact to efficiently estimate the transition and reward tensors in
discrete multi-agent reinforcement learning problems.

###  2.2 Reinforcement learning

In reinforcement learning, the goal is to compute a strategy to perform a
certain task. Tasks are formalized as Markov decision processes (MDPs)
$(S,A,T,R,\gamma)$, where $S$ and $A$ are the state and action spaces, $T$ and
$R$ are the transition and reward function and $\gamma$ is the discount
factor. At each time step $t\in\mathbb{N}$ an agent chooses an action
$a_{t}\in A$ based on the state $s_{t}\in S$. The environment then returns a
reward $r_{t}=R(s_{t},a_{t})$ and the next state $s_{t+1}\sim
T(\cdot|s_{t},a_{t})$. The strategy to choose actions is called the policy
$\pi:S\times A\rightarrow[0,1]$ which defines a probability distribution over
the actions given the current state. The goal of the agent is to maximize the
expected discounted reward
$\mathbb{E}_{T,\pi}\left[\sum_{t=1}^{H}\gamma^{t}r_{t}\right]$, where the
expectation is over the states and actions, whose probability distributions
are implied by the transition function $T$ and policy $\pi$.

In multi-agent reinforcement learning (MARL), there are multiple agents that
interact with the environment as opposed to only one agent. Each agent has its
own action space $A_{i}$, which means that the transition function is now a
function $S\times A_{1}\times\cdots\times A_{n}\times S\rightarrow[0,1]$ and
the reward function is $S\times A_{1}\times A_{2}\times\dots\times
A_{n}\rightarrow[0,1]$.

Clearly this can be cast as a single agent reinforcement learning problem by
setting $A=A_{1}\times\cdots\times A_{n}$. A result of this is that the action
space grows exponentially large with the number of agents, further increasing
the requirement of efficient exploration. Instead of casting it as a single
agent reinforcement learning problem, explicitly incorporating the multi-agent
paradigm allows to exploit more structure in the MDP.

In this work, this is done by considering the transition function to be a
tensor $T\in\mathbb{R}^{S\times A_{1}\times\cdots\times A_{n}\times S}$ such
that $T_{sa_{1}\dots a_{n}s^{\prime}}=T(s,a_{1},\dots,a_{n},s^{\prime})$. The
reward function is analogously written as a tensor $R\in\mathbb{R}^{s\times
a_{1}\times\cdots\times a_{n}}$. If the tensors $T$ and $R$ are of low rank,
models formed by an agent during training can be expected to generalize across
unseen state-action pairs.

##  3 Methods

###  3.1 Tensor decomposition algorithms

While computing tensor decompositions is NP-hard in general Hillar and Lim
[2013], there exist algorithms such as Harshman [1970] and Anandkumar et al.
[2015] which are proven to converge in special cases. The algorithm used in
this work is an ablation of the alternating rank 1 updates algorithm presented
in Anandkumar et al. [2015]. The restarts, clustering and clipping procedures
are left out to form a shorter and simpler algorithm which still performs well
in practice. The algorithm as used is presented in algorithms 3, 4, and 1. The
main idea of the algorithm is to run asymmetric power updates to compute a
good starting value for alternating minimization, which further improves the
accuracy of the decomposition.

Input: A tensor $T\in\mathbb{R}^{m\times n\times p}$ and decomposition rank
$r$

Result: $\\{u_{k}^{j}\\}_{k=1\dots r}^{j=1\dots n},\\{w_{k}\\}_{k=1\dots r}$
such that $T\approx\sum_{k=1}^{r}w_{k}u_{k}^{1}\otimes
u_{k}^{2}\otimes\cdots\otimes u_{k}^{n}$

$(\\{u^{j}_{k}\\}_{k=1\dots r}^{j=1\dots n},\\{w_{k}\\}_{k=1\dots
r})=\text{PowerIteration}(T,r)$ (algorithm 3);

$(\\{u^{j}_{k}\\}_{k=1\dots r}^{j=1\dots n},\\{w_{k}\\}_{k=1\dots
r})=\text{AlternatingMinimize}(T,\\{u^{j}_{k}\\}_{k=1\dots r}^{j=1\dots
n},\\{w_{k}\\}_{k=1\dots r})$ (algorithm 4);

Algorithm 1 Alternating rank 1 updates

###  3.2 Tensor completion

In the tensor completion problem, the goal is to recover a tensor with only
partially observed entries. For a tensor
$T\in\mathbb{R}^{d_{1}\times\cdots\times d_{n}}$, let
$\Omega\in\\{0,1\\}^{d_{1}\times\cdots\times d_{n}}$ denote a mask such that
$\Omega_{i_{1}\dots i_{n}}=1$ if and only if entry $T_{i_{1}\dots i_{n}}$ has
been observed. A method proposed in Jain and Oh [2014] and Liu and Moitra
[2020] involves solving the minimization problem

$\operatorname*{argmin}_{\\{u_{k}^{j}\\}_{k=1\dots r}^{j=1\dots
n}\in\mathbb{R}^{d_{j}},\\{w_{k}\\}_{k=1\dots r}\in\mathbb{R}}\|\Omega\cdot
T-\Omega\cdot\sum_{k=1}^{r}w_{k}u_{k}^{1}\otimes u_{k}^{2}\otimes\cdots\otimes
u_{k}^{n}\|_{F},$

where $\cdot$ denotes an entrywise multiplication. Algorithm 1 can be used to
solve this problem with a slight modification to the alternating minimization
step as showcased in appendix D.

###  3.3 Model based reinforcement learning

In model based reinforcement learning agents make models of the environment to
plan ahead, instead of attempting to maximize reward directly. Under the
assumption that the transition and reward tensors are of low rank, using
tensor decomposition allows for sample efficient models that generalize over
unseen state-action pairs.

The deterministic reward tensor is estimated using tensor completion, where
the unobserved entries are simply the state-action pairs the agents have never
experienced. After enough exploration, enough entries of the reward tensor
will be revealed to reconstruct the entire tensor.

The algorithm presented in this paper follows a very standard model-based
reinforcement learning approach and is presented in algorithm 2. The NORMALIZE
function is an entry wise division so that the resulting tensor is a
transition tensor, that is the sum over the resulting states is 1. The
POLICYIMPROVEMENT function is clarified in appendix C.

###  3.4 Relationship to Tesseract

This method differs from Mahajan et al. [2021] because in this work,
decompositions of the entire tensor $T\in\mathbb{R}^{S\times
A_{1}\times\cdots\times A_{n}\times S}$ and $R\in\mathbb{R}^{S\times
A_{1}\times\cdots\times A_{n}}$ are computed. Model based Tesseract instead
considers for each $s,s^{\prime}\in S$ the $A_{1}\times\cdots\times A_{2}$
tensor $\tilde{T}_{ss^{\prime}}=T(e_{s},I,\dots,I,e_{s}^{\prime})$, and
computes an individual tensor decomposition for each state and next state
pair. Analogously, it considers for each state $s\in S$ the reward tensor
$R_{s}=R(e_{s},I,\dots,I)$ and computes a decomposition for every state.

In theory, both methods can represent the same transition and reward
functions. To see this, consider for example an MDP with 2 agents, such that
the reward tensor is of order 3. Let $R=\sum_{i=1}^{r}w_{i}x_{i}\otimes
y_{i}\otimes z_{i}$ be the true reward tensor. This can represented by
Tesseract by setting $R_{s}=\sum_{i=1}^{r}w_{i}\langle e_{s},x_{i}\rangle
y_{i}\otimes z_{i}$, where $e_{s}$ is the $s$-th standard basis vector. This
results in the combined reward tensor

$\sum_{s\in S}e_{s}\otimes R_{s}=\sum_{s\in
S}e_{s}\otimes\sum_{i=1}^{r}w_{i}\langle e_{s},x_{i}\rangle y_{i}\otimes
z_{i}=\sum_{i=1}^{r}w_{i}\sum_{s\in S}e_{s}\otimes\langle e_{s},x_{i}\rangle
y_{i}\otimes z_{i}=R$

Conversely, if $R$ is of the form $\sum_{s\in S}e_{s}\otimes R_{s}$ where each
$R_{s}$ is of rank $r$, then the rank of $R$ is bounded by $|S|r$ so it can be
represented in our framework. Thus, when low rank structure spans across
states, out method would ensure better sample efficiency as it would require
fewer number of parameters.

Input: An MDP, state space size $S$, action space sizes $A_{1},\dots A_{n}$,
Hyperparameters:
$n_{\texttt{episodes}},n_{\texttt{train}},\epsilon,n_{\texttt{improvement
iter}}$

Result: Policy $\pi$ with good performance

$D=0\in\mathbb{N}^{S\times A_{1}\times\cdots\times A_{n}\times S}$;

$R=0\in\mathbb{R}^{S\times A_{1}\times\cdots\times A_{n}}$;

Initialize random $\pi$;

for  _$\texttt{episode}=1,2,\dots,n_{\texttt{episodes}}$_ do

$s_{1}\sim MDP$;

$a_{1}\sim\begin{cases}\pi(s_{0})&\text{ with probability
}1-\epsilon(\texttt{episode})\\\ \texttt{Unif}(A)&\text{ with probability
}\epsilon(\texttt{episode})\end{cases}$;

for  _$t=1,\dots,\texttt{episode length}$_ do

$(s_{t+1},r_{t})\sim MDP(.|s_{t},a_{t})$;

$D_{s_{t}a_{t}s_{t+1}}=D_{s_{t}a_{t}s_{t+1}}+1$;

$R_{s_{t}a_{t}}=r_{t}$;

$a_{t}\sim\begin{cases}\pi(s_{t})&\text{ with probability
}1-\epsilon(\texttt{episode})\\\ \texttt{Unif}(A)&\text{ with probability
}\epsilon(\texttt{episode})\end{cases}$;

if  _$n_{\texttt{episodes}}\,\%\,n_{\texttt{train}}=0$_ then

$\hat{T}=\texttt{DECOMP}(\texttt{NORMALIZE}(D),r_{T})$;

$\Omega=D>0$;

$\hat{R}=\texttt{TENSORCOMPLETION}(R,\Omega,r_{R})$;

$\pi=\texttt{POLICYIMPROVEMENT}(\pi,\hat{T},\hat{R},n_{\texttt{improvement
iter}})$;

Algorithm 2 CP-Decomposed state-action space reinforcement learning

##  4 Experiments

###  4.1 Random transition and reward functions of predefined rank

This experiment involves algorithm 2 applied to an MDP described by a randomly
generated transition tensor $T$ and reward tensor $R$. The MDP has 20 states,
and 3 agents with 10 actions each, leading to 20000 state-action pairs. Both
the transition and reward tensors are of rank $5$. More information on how
they are generated can be found in appendix E.

We tested three different agents for experiments. The first agent is a
baseline agent which uses no decompositions. It uses the maximum likelihood
estimator for $T$ and fills in missing rewards for unvisited state-action
pairs with the mean of the visited rewards. Secondly, an agent using
decomposition across entire state action space for for $T$ and $R$ with three
settings of approximate rank $5$ (exact), $3$ (insufficient), $10$
(overparametrised). The final agent is model based Tesseract with rank $5$ and
$1$ decompositions. Note that the rank $5$ case can represent the correct
transition and reward tensors, but is overparametrised for the task (60000
parameters versus 350 for our rank $5$ agent for the transition function).
Similarly, the rank $1$ agent will be insufficient for representing the actual
dynamics but will provide faster learning.

For each agent, if the slice through the estimated transition tensor $T$
corresponding to a specific state-action pair contains only zeros, all entries
are set to $\frac{1}{|S|}$. This means that if there is no estimate for
$T(\cdot|s,a_{1},a_{2},a_{3})$, a uniform distribution is assumed instead.
This happens for the no decompositions agent exactly when a state-action pair
has never been visited before.

Each agent is trained for 200 episodes, recomputing their models and applying
policy improvement every 10 episodes. The agents use $\epsilon$ greedy
exploration with epsilon decaying from $0.9$ to $0.1$. During training, the
total episodic rewards, errors in the transition tensor and errors in the
reward tensors are tracked. The entire experiment is ran 20 times, with newly
generated $T$ and $R$ for each run. The optimal reward in each experiment is
computed beforehand via policy improvement on the true functions $T$ and $R$,
and then for each experiment the optimal reward is subtracted from the
episodic rewards so that optimal performance is a reward of $0$ for each
experiment. Finally, the number of unique visited state-action pairs is also
tracked. The results of the experiment are shown in figures 1, 2 and 5(a).
Figure 1 shows that algorithm 2 significantly outperforms a standard model-
based approach without tensor decompositions in the setting where $T$ and $R$
are of low rank. While the agent with rank 3 decompositions achieves a
slightly sub-optimal policy, the performance seems to be quite robust against
incorrectly guessing the correct rank for the problem.

Interestingly, agents without tensor decompositions outperform the agents that
use tensor decompositions during the first few episodes. This can be
attributed to the unrobustness of tensor completion. Figure 1 shows that for
our algorithm, during the first 20 episodes the error in the reward tensor can
be of order $10^{5}$ and higher, because the optimization problem is very ill-
conditioned when little entries are revealed. Tesseract suffers even more from
this problem, as each individual state now requires sufficiently many revealed
entries. A way to overcome these problems could be for example to take the
naive estimate without tensor completion when attempting tensor completion
results in very extreme values, or adding regularization to the optimization
problem. Figure 1 also shows that with sufficiently many revealed entries, our
method achieves very good approximates of the reward tensor. If the
approximate reward tensor rank is set correctly (rank 5), the reward tensor is
recovered almost exactly after visiting only 4000 (see figure 5(a)) or 20% of
the state-action pairs. Setting the rank results in slower convergence, but
still yields a reasonably good estimate. Finally, setting the rank too low
causes the agent to be incapable of representing the true reward tensor, but
on limited revealed entries this estimate still outperforms the estimate
without decompositions. Tesseract with rank 5 decompositions takes a long time
to get a good estimate, but eventually outperforms our method with rank 3
decompositions. This is explained by the analysis in 3.4, which showed that
rank 5 Tesseract is in theory capable to represent the true reward tensor,
albeit requiring many more samples in comparison as confirmed by this
experiment.

Figure 2 shows the error in transition tensors. Note that the error of
Tesseract and the agent without decompositions increases over time. This is
due to the fact that for many states-action pairs, the default uniform
distribution assigning probability $\frac{1}{|S|}$ to each state is a better
estimate than an extreme distribution resulting from only one observation of
that state-action pair. Figure 5(a) shows that even after 200 episodes, only
around 11000 state action pairs out of 20000 total are visited, meaning that
many state-action pairs are likely to have been visited only once. This means
that unless an agent can combine information from different state-action
pairs, it is unfeasible to make a good transition function estimate. Since
using no decompositions assumes every state-action pair to be independent,
there is no generalization across states-action pairs. Tesseract does slightly
better as it attempts to generalize the action space for each state
independently, but figure 2 shows that our method produces significantly
better transition tensor estimates by attempting to generalize over the
combined state-action space.

Figure 1: Total reward per episode at test time (left) and sum of squared
errors of the reward tensor estimate (right) after a number of episodes of
training. Figure 2: Mean squared error of the estimate of the transition
tensor after training a number of episodes in the first experiment. The right
plot is a zoomed in version of the left plot.

###  4.2 MDP with degenerate states

In this experiment we test algorithm 2 on state degeneracy, a situation where
our method can provide further sample efficiency. State degeneracy can occur
when observations are noisy. We use an MDP with 3 agents, this time with 16
states and each agent has an action space of size 20. The 16 states are split
into 4 groups, where each group has the same transition function of rank 1,
and a linearly dependent reward tensor of rank 1. This means that rank 1
Tesseract is expected to be able to recover exact models after enough
iterations. Furthermore, similar to the analysis in 3.4, writing
$R=\sum_{i=1,5,9,13}(e_{i}+e_{i+1}+e_{i+2}+e_{i+3})\otimes R_{i}$ where
$R_{i}$ denotes a reward tensor for each group, reveals that the entire reward
tensor is of rank at most 4. Similar to the previous experiment, we consider
the following agents: A baseline agent using no decomposition, Agents using
decomposition across state action space with ranks $4$ and $8$, Tesseract with
rank $4$ and $1$. The entire experiment is repeated 20 times. The results are
shown in figures 3, 4 and 5(b). Like in the results of the first experiment,
figure 3 shows that the agents that use decompositions accross state-action
space outperform the other agents in terms of total reward obtained. This is
mostly attributable to the performance on the reward tensor error. We also
observe that Tesseract is unable to recover a good reward tensor estimate in
the given sample budget. In contrast to the previous experiment, the
transition tensor estimates of the agents that use state-action decompositions
do not differ significantly from the estimates made by Tesseract. This can be
explained by the fact that in this experiment, the groups themselves have
entirely independent transition and reward functions, which means that
generalization is only possible within groups. In contrast, the low rank
structure imposed on the entire transition tensor in the first experiment
allowed our method to generalize over all states.

Figure 3: Total reward per episode at test time (left) and sum of squared
errors of the reward tensor estimate (right) after a number of episodes of
training in the state degeneracy MDP.

##  5 Conclusion

In this position paper we investigated whether tensor decompositions can be
used across state actions space for better sample efficiency in RL for the
model based setting. Our experiments show that an algorithm which computes CP-
decompositions of the environment models has significant advantages when the
MDP is described by low rank transition and reward functions.

## References

  * Mahajan et al. [2021] Anuj Mahajan, Mikayel Samvelyan, Lei Mao, Viktor Makoviychuk, Animesh Garg, Jean Kossaifi, Shimon Whiteson, Yuke Zhu, and Animashree Anandkumar.  Tesseract: Tensorised actors for multi-agent reinforcement learning.  In _Proceedings of the 38th International Conference on Machine Learning_ , volume 139, pages 7301–7312. PMLR, 2021.  URL https://proceedings.mlr.press/v139/mahajan21a.html.
  * DeepMind-OEL et al. [2021] DeepMind-OEL, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, Nat McAleese, Nathalie Bradley-Schmieg, Nathaniel Wong, Nicolas Porcel, Roberta Raileanu, Steph Hughes-Fitt, Valentin Dalibard, and Wojciech Marian Czarnecki.  Open-ended learning leads to generally capable agents.  _arXiv preprint arXiv:2107.12808_ , 2021.
  * Mahajan et al. [2019] Anuj Mahajan, Tabish Rashid, Mikayel Samvelyan, and Shimon Whiteson.  Maven: Multi-agent variational exploration.  In _Advances in Neural Information Processing Systems_ , pages 7611–7622, 2019.
  * Gupta et al. [2020] Tarun Gupta, Anuj Mahajan, Bei Peng, Wendelin Böhmer, and Shimon Whiteson.  Uneven: Universal value exploration for multi-agent reinforcement learning.  _arXiv preprint arXiv:2010.02974_ , 2020.
  * Wang et al. [2020a] Tonghan Wang, Tarun Gupta, Anuj Mahajan, Bei Peng, Shimon Whiteson, and Chongjie Zhang.  Rode: Learning roles to decompose multi-agent tasks.  _arXiv preprint arXiv:2010.01523_ , 2020a.
  * Wang et al. [2020b] Jianhao Wang, Zhizhou Ren, Terry Liu, Yang Yu, and Chongjie Zhang.  Qplex: Duplex dueling multi-agent q-learning.  _arXiv preprint arXiv:2008.01062_ , 2020b.
  * Hillar and Lim [2013] Christopher J. Hillar and Lek-Heng Lim.  Most tensor problems are np-hard.  _J. ACM_ , 60(6), November 2013.  ISSN 0004-5411.  doi: 10.1145/2512329.  URL https://doi.org/10.1145/2512329.
  * Harshman [1970] Richard Harshman.  Foundations of the parafac procedure: Models and conditions for an "explanatory" multi-modal factor analysis.  _UCLA Working Papers in Phonetics_ , 16, 1970.
  * Anandkumar et al. [2015] Animashree Anandkumar, Rong Ge, and Majid Janzamin.  Guaranteed non-orthogonal tensor decomposition via alternating rank-$1$ updates, 2015.
  * Jain and Oh [2014] Prateek Jain and Sewoong Oh.  Provable tensor factorization with missing data.  In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger, editors, _Advances in Neural Information Processing Systems_ , volume 27. Curran Associates, Inc., 2014.  URL https://proceedings.neurips.cc/paper/2014/file/c15da1f2b5e5ed6e6837a3802f0d1593-Paper.pdf.
  * Liu and Moitra [2020] Allen Liu and Ankur Moitra.  Tensor completion made practical.  In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 18905–18916. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/dab1263d1e6a88c9ba5e7e294def5e8b-Paper.pdf.
  * Anandkumar et al. [2012] Animashree Anandkumar, Daniel Hsu, and Sham M. Kakade.  A method of moments for mixture models and hidden markov models.  In Shie Mannor, Nathan Srebro, and Robert C. Williamson, editors, _Proceedings of the 25th Annual Conference on Learning Theory_ , volume 23 of _Proceedings of Machine Learning Research_ , pages 33.1–33.34, Edinburgh, Scotland, 25–27 Jun 2012. PMLR.  URL https://proceedings.mlr.press/v23/anandkumar12.html.
  * Anandkumar et al. [2014] Animashree Anandkumar, Rong Ge, Daniel Hsu, Sham M. Kakade, and Matus Telgarsky.  Tensor decompositions for learning latent variable models.  _J. Mach. Learn. Res._ , 15(1):2773–2832, January 2014.  ISSN 1532-4435.
  * Cichocki et al. [2017] Andrzej Cichocki, A. Phan, Qibin Zhao, Namgil Lee, I. Oseledets, Masashi Sugiyama, and Danilo P. Mandic.  Tensor networks for dimensionality reduction and large-scale optimization: Part 2 applications and future perspectives.  _Found. Trends Mach. Learn._ , 9:431–673, 2017.
  * Cheng et al. [2017] Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang.  A survey of model compression and acceleration for deep neural networks.  _CoRR_ , abs/1710.09282, 2017.  URL http://arxiv.org/abs/1710.09282.
  * Kossaifi et al. [2019] Jean Kossaifi, Adrian Bulat, Georgios Tzimiropoulos, and Maja Pantic.  T-net: Parametrizing fully convolutional nets with a single high-order tensor.  _2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 7814–7823, 2019.
  * Kossaifi et al. [2020] Jean Kossaifi, Antoine Toisoul, Adrian Bulat, Yannis Panagakis, Timothy M. Hospedales, and Maja Pantic.  Factorized higher-order cnns with an application to spatio-temporal emotion estimation.  _2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 6059–6068, 2020.
  * Bulat et al. [2020] Adrian Bulat, Jean Kossaifi, Georgios Tzimiropoulos, and Maja Pantic.  Incremental multi-domain learning with network latent tensor factorization.  In _AAAI_ , 2020.
  * Sunehag et al. [2018] Peter Sunehag, Guy Lever, Audrunas Gruslys, Wojciech Marian Czarnecki, Vinicius Zambaldi, Max Jaderberg, Marc Lanctot, Nicolas Sonnerat, Joel Z. Leibo, Karl Tuyls, and Thore Graepel.  Value-decomposition networks for cooperative multi-agent learning based on team reward.  In _Proceedings of the 17th International Conference on Autonomous Agents and MultiAgent Systems_ , AAMAS ’18, page 2085–2087, Richland, SC, 2018. International Foundation for Autonomous Agents and Multiagent Systems.
  * Rashid et al. [2018] Tabish Rashid, Mikayel Samvelyan, Christian Schroeder, Gregory Farquhar, Jakob Foerster, and Shimon Whiteson.  QMIX: Monotonic value function factorisation for deep multi-agent reinforcement learning.  In Jennifer Dy and Andreas Krause, editors, _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pages 4295–4304. PMLR, 10–15 Jul 2018.  URL https://proceedings.mlr.press/v80/rashid18a.html.
  * Bromuri [2012] Stefano Bromuri.  A tensor factorization approach to generalization in multi-agent reinforcement learning.  In _2012 IEEE/WIC/ACM International Conferences on Web Intelligence and Intelligent Agent Technology_ , volume 2, pages 274–281, 2012\.  doi: 10.1109/WI-IAT.2012.21.
  * Azizzadenesheli et al. [2016] Kamyar Azizzadenesheli, Alessandro Lazaric, and Animashree Anandkumar.  Reinforcement learning of pomdps using spectral methods.  06 2016.

##  Appendix A Related work

Work on tensor decompositions in general machine learning include Anandkumar
et al. [2012], which uses the CP-decomposition to learn mixture models and
hidden Markov models, and Anandkumar et al. [2014] learns latent variable
models. Another application of tensor methods is to compress neural networks
in Cichocki et al. [2017] and Cheng et al. [2017]. Kossaifi et al. [2019]
parametrizes convolutional nets with a high-order tensor of low rank to reduce
over-parameterization, with applications to spatio-temporal tasks in Kossaifi
et al. [2020]. In Bulat et al. [2020] this parametrization is used for multi-
domain image classification.

Previous reinforcement learning methods that attempt to exploit structure in
the multi-agent setting include VDN Sunehag et al. [2018], which models the
joint Q-function (see appendix C) as a sum of the agents individual
Q-functions. This is generalised by QMIX Rashid et al. [2018], which learns a
monotonic function of the individual Q-functions instead of taking a sum.

Methods for generalization in multi agent reinforcement learning using
specifically tensor decompositions include Bromuri [2012], where tensor
decompositions are used to factorize the $Q$-function in model free
$Q-$learning algorithms. Our method instead factorizes the estimated
transition and reward functions of the MDP in a model based algorithm. This
idea was initially proposed in Mahajan et al. [2021], which contains a model
free algorithm and a model based algorithm. The difference between the model
based algorithm in Mahajan et al. [2021] and our method, is that we factorize
over the state-action space, while Mahajan et al. [2021] factorizes only over
the action space. This allows our work to potentially generalize over unseen
states instead of only over unseen actions.

Work in tensor decompositions for partially observable MDPs (POMPDs) in a
single agent setting include Azizzadenesheli et al. [2016]. Adapting our
method for generalisation in multi agent MDPs for POMPDs is an interesting
future research direction

##  Appendix B Decomposition algorithms

Input: A tensor $T\in\mathbb{R}^{d_{1}\times\cdots\times d_{n}}$ and
decomposition rank $r$, tolerance $\epsilon$

Result: Vectors $\\{u_{k}^{j}\\}_{k=1\dots r}^{j=1\dots n}$ and scalars
$\\{w_{k}\\}_{k=1\dots r}$ such that
$T\approx\sum_{k=1}^{r}w_{k}u_{k}^{1}\otimes u_{k}^{2}\otimes\cdots\otimes
u_{k}^{n}$

for  _k = 1, …r_ do

for  _j = 1, …, n_ do

$u_{k,0}^{j}\sim\mathcal{N}(0,1)$;

$u_{k,0}^{j}=\frac{u_{k}^{j}}{\|u_{k}^{j}\|}$;

end for

while  _$\sum_{j=1}^{n}\|u_{k,n+1}^{j}-u^{j}_{k,n}\|^{2} >\epsilon$_ do

for  _j=1, …, n_ do

$u^{j}_{k,m+1}=\frac{T(u^{1}_{k,m},\dots,u^{j-1}_{k,m},I,u^{j+1}_{k,m},\dots,u^{n}_{k,m})}{\|T(u^{1}_{k,m},\dots,u^{j-1}_{k,m},I,u^{j+1}_{k,m},\dots,u^{n}_{k,m})\|}$;

end for

end while

$w_{k}=T(u_{k}^{1},\dots,u_{k}^{n})$;

$w_{k}^{N}=c_{k}w_{k}^{N}$;

$T=T-w_{k}u_{k}^{1}\otimes u_{k}^{2}\otimes\cdots\otimes u_{k}^{n}$;

end for

Algorithm 3 Tensor power iteration with deflation

Input: A tensor $T$, starting values for $\\{u_{k}^{j}\\}_{k=1\dots
r}^{j=1\dots n},\\{w_{k}\\}_{k=1\dots r}$

Result: $\\{u_{k}^{j}\\}_{k=1\dots r}^{j=1\dots n},\\{w_{k}\\}_{k=1\dots r}$
such that $T\approx\sum_{k=1}^{r}w_{k}u_{k}^{1}\otimes
u_{k}^{2}\otimes\cdots\otimes u_{k}^{n}$

while  _stopping criterion_ do

for  _$l=1,\dots,r$_ do

for  _$j=1,\dots,n$_ do

$u_{l}^{j}={(T-\sum_{k\neq l}^{r}u_{k}^{1}\otimes
u_{k}^{2}\otimes\cdots\otimes
u_{k}^{n})(u_{l}^{1},\dots,u_{l}^{j-1},I,u_{l}^{j+1},\dots,u_{l}^{n})}$;

$u_{l}^{j}=\frac{u_{l}^{j}}{\|u_{l}^{j}\|}$;

end for

$w_{l}=(T-\sum_{k\neq l}^{r}u_{k}^{1}\otimes u_{k}^{2}\otimes\cdots\otimes
u_{k}^{n})(u_{l}^{1},\dots,u_{l}^{n})$;

end for

end while

Algorithm 4 Alternating minimization

##  Appendix C Policy improvement

Policy improvement is a well known method in reinforcement learning to compute
optimal policies with respect to the MDP parameters $T$ and $R$. In model
based reinforcement learning algorithms, this is used to compute a good policy
after estimating $T$ and $R$ with estimates $\hat{T}$ and $\hat{R}$. If the
estimates are close enough, the optimal policy with respect to $\hat{T}$,
$\hat{R}$ will also perform well on the actual MDP described by $T,R$.

Computing the optimal policy uses the $Q$-function, which maps each state $s$
and each action $a$ to the expected reward of executing action $a$ in state
$s$ and following policy $\pi$ afterwards. This can recursively be written as

$Q^{\pi}(s,a)=R(s,a)+\gamma\sum_{s^{\prime}\in
S}T(s^{\prime}|a,s)\sum_{a^{\prime}\in
A}\pi(a^{\prime}|s^{\prime})Q^{\pi}(s^{\prime},a^{\prime})$

The value function maps each state to the expected future reward in the state
when following policy $\pi$. This can be computed from the $Q$-function via

$V^{\pi}(s)=\sum_{a\in A}Q(s,a)\pi(a|s)$

The main idea of policy improvement is to iteratively select states for which
$V^{\pi}(s)<\max_{a\in A}Q^{\pi}(s,a)$. This means that there exists an action
which achieves better reward than the current policy, so the policy $\pi$ is
updated to use the better action instead. After this, since the policy has
changed, $Q^{\pi}$ and $V^{\pi}$ need to be computed again to repeat this
process. This is guaranteed to converge to an optimal policy eventually on
finite state and action spaces.

Input: Starting policy $\pi_{0}$

Result: Improved policy $\pi$

Set $i=0$;

Compute $Q^{\pi_{i}}$ and $V^{\pi_{i}}$;

while  _there exist $s\in S$ such that $V^{\pi_{i}}(s)<\max_{a}Q^{\pi}(s,a)$_
do

Pick $s$ such that $V^{\pi_{i}}(s)<\max_{a}Q^{\pi_{i}}(s,a)$;

Let $A^{*}_{s,\pi_{i}}=\operatorname*{argmax}_{a}Q^{\pi}(s,a)$;

Define a new policy: $\pi_{i+1}=\pi_{i}$;

Set $\pi_{i+1}(A^{*}_{s,\pi_{i}}|s)=1$;

$i=i+1$;

Compute $Q^{\pi_{i}}$ and $V^{\pi_{i}}$;

end while

Algorithm 5 Policy improvement

##  Appendix D Tensor completion

To modify algorithm 1 for tensor completion, alternating minimization
(algorithm 4) can be modified to solve

$\operatorname*{argmin}_{u_{l}^{j}\in\mathbb{R}^{d_{j}}}\|\Omega\cdot
T-\Omega\cdot\sum_{k=1}^{r}w_{k}u_{k}^{1}\otimes u_{k}^{2}\otimes\cdots\otimes
u_{k}^{n}\|_{F}$

at each iteration, instead of the usual problem

$\operatorname*{argmin}_{u_{l}^{j}\in\mathbb{R}^{d_{j}}}\|T-\sum_{k=1}^{r}w_{k}u_{k}^{1}\otimes
u_{k}^{2}\otimes\cdots\otimes u_{k}^{n}\|_{F}.$

This leads to the update

$u_{l}^{j}=\frac{(\Omega\cdot T-\Omega\cdot\sum_{k\neq
l}^{r}w_{k}u_{k}^{1}\otimes u_{k}^{2}\otimes\cdots\otimes
u_{k}^{n})(u_{l}^{1},\dots,u_{l}^{j-1},I,u_{l}^{j+1},\dots,u_{l}^{n})}{\Omega(u_{l}^{1}\cdot
u_{l}^{1},\dots,u_{l}^{j-1}\cdot u_{l}^{j-1},I,u_{l}^{j+1}\cdot
u_{l}^{j+1},\dots,u_{l}^{n}\cdot u_{l}^{n})}$

instead of the usual update displayed in algorithm 4. Similarly, the update
for the weights is given by

$w_{l}=\frac{(\Omega\cdot T-\Omega\cdot\sum_{k\neq l}^{r}w_{k}u_{k}^{1}\otimes
u_{k}^{2}\otimes\cdots\otimes
u_{k}^{n})(u_{l}^{1},\dots,u_{l}^{n})}{\Omega(u_{l}^{1}\cdot
u_{l}^{1},\dots,u_{l}^{n}\cdot u_{l}^{n})}.$

This modification is inspired by the tensor completion method in Jain and Oh
[2014], which is proven to work for symmetric orthogonal tensors.

##  Appendix E Tensor generation details

The target reward and transition tensors in the experiments are generated by
algorithm 6. The weights are chosen to be $w=\texttt{linspace}(0.1,1)$ for the
reward tensor in the first experiment.

Input: Desired shape $(d_{1},\dots,d_{n})$ and rank $r$, weights
$w\in\mathbb{R}^{r}$

Result: Rank $r$ tensor with dimensions $d_{1},\dots,d_{n}$

for  _$i=1,\dots,n$_ do

for  _$l=1,\dots,r$_ do

$u_{i}^{l}\sim\mathcal{N}(0,I_{d_{i}})$;

$u_{i}^{l}=\frac{u_{i}^{l}}{\|u_{i}^{l}\|}$;

$T=\sum_{l=1}^{r}w_{l}u_{1}^{l}\otimes\cdots\otimes u_{n}^{l}$;

Algorithm 6 Tensor generation

A complication is that the transition tensor must satisfy
$\sum_{s^{\prime}}T_{s}a_{1}a_{2}a_{3}s^{\prime}=1$. It is difficult to
directly generate a tensor of fixed rank with this property, and normalizing a
tensor by setting

$T^{\prime}_{s_{1}a_{1}a_{2}a_{3}s_{2}}=\frac{T_{s_{1}a_{1}a_{2}a_{3}s_{2}}}{\sum_{s^{\prime}}T_{s_{1}a_{1}a_{2}a_{3}s^{\prime}}}$

changes the rank of the tensor. This is overcome by iteratively normalizing,
computing a new decomposition of higher than the desired rank, and then
truncating it to the desired rank. Repeating this as shown in algorithm 7
appears to converge in practice, enabling the generation of a valid transition
tensor that is arbitrarily close to to a tensor of desired rank.

Input: Desired shape $(d_{1},\dots,d_{n})$ and rank $r$, tolerance $\epsilon$

Result: Approximately rank $r$ transition tensor with dimensions
$d_{1},\dots,d_{n}$

$T=\texttt{GenerateTensor}((d_{1},\dots,d_{n}),r,w)$;

while  _ $\|\texttt{Normalize}(T)-T\|_{F}>\epsilon$_ do

$T=\texttt{Normalize(T)}$;

$T=\texttt{TensorDecomp(T, 2r)}$;

$T=\texttt{Truncate}(T,r)$;

$T=\texttt{Normalize}(T)$;

Algorithm 7 Transition tensor generation

##  Appendix F Extra experiment figures

Figure 4: Mean squared error of the estimate of the transition tensor after
training a number of episodes in the state degeneracy MDP. The right plot is a
zoomed in version of the left plot.

(a) Number of unique states visited after training for a number of episodes in
the first experiment. The MDP has 20000 unique states in total

(b) Number of unique states visited after training for a number of episodes in
the state degeneracy experiment. The total number of state-action pairs in the
MDP is 128000