Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning

## Abstract

Abstract Effective decision-making in partially observable environments demands
robust memory management. Despite their success in supervised learning,
current deep-learning memory models struggle in reinforcement learning
environments that are partially observable and long-term. They fail
to efficiently capture relevant past information, adapt flexibly to
changing observations, and maintain stable updates over long episodes.
We theoretically analyze the limitations of existing memory models
within a unified framework and introduce the Stable Hadamard Memory,
a novel memory model for reinforcement learning agents. Our model
dynamically adjusts memory by erasing no longer needed experiences
and reinforcing crucial ones computationally efficiently. To this
end, we leverage the Hadamard product for calibrating and updating
memory, specifically designed to enhance memory capacity while mitigating
numerical and learning challenges. Our approach significantly outperforms
state-of-the-art memory-based methods on challenging partially observable
benchmarks, such as meta-reinforcement learning, long-horizon credit
assignment, and POPGym, demonstrating superior performance in handling
long-term and evolving contexts.

## 1 Introduction

Reinforcement learning agents necessitate memory. This is especially
true in Partially Observable Markov Decision Processes (POMDPs Kaelbling et al. ()),
where past information is crucial for making informed decisions. However,
designing a robust memory remains an enduring challenge, as agents
must not only store long-term memories but also dynamically update
them in response to evolving environments. Memory-augmented neural
networks (MANNs)–particularly those developed for supervised learning
Graves et al. (); Vaswani et al. (), while offering promising
solutions, have consistently struggled in these dynamic settings.
Recent empirical studies Morad et al. (); Ni et al. ()
have shown that MANNs exhibit instability and underperform simpler
vector-based memory models such as GRU Chung et al. ()
or LSTM Hochreiter (). The issue is exacerbated in complex
and sparse reward scenarios where agents must selectively retain and
erase memories based on relevance. Unfortunately, existing methods
fail to provide a memory writing mechanism that is simultaneously
efficient, stable, and flexible to meet these demands.

In this paper, we focus on designing a better writing mechanism to
encode new information into the memory*.* To this end, we introduce
the Hadamard Memory Framework (HMF), a unified model that encompasses
many existing writing methods as specific cases. This framework highlights
the critical role of *memory calibration*, which involves linearly
adjusting memory elements by multiplying the memory matrix with a
*calibration matrix* and then adding an* update matrix*.
By leveraging Hadamard products that operate element-wise on memory
matrices, we allow memory writing without mixing the memory cells
in a computationally efficient manner. More importantly, the calibration
and update matrices are dynamically computed based on the input at
each step. This enables the model to learn adaptive memory rules,
which are crucial for generalization. For instance, in meta-reinforcement
learning with varying maze layouts, a fixed memory update rule may
work for one layout but fail in another. By allowing the calibration
matrix to adjust according to the current maze observation, the agent
can learn to adapt to any layout configuration. A dynamic calibration
matrix also enables the agent to flexibly forget and later recall
information as needed. For example, an agent navigating a room may
need to remember the key’s location, retain it during a detour, and
later recall it when reaching a door, while discarding irrelevant
detour events.

Although most current memory models can be reformulated within the
HMF, they tend to be either overly simplistic with limited calibration
capabilities Katharopoulos et al. (); Radford et al. ()
or unable to manage memory writing reliably, suffering from gradient
vanishing or exploding issues Ba et al. (); Morad et al. ().
To address these limitations, we propose a specific instance of HMF,
called Stable Hadamard Memory, which introduces a novel calibration
mechanism based on two key principles: (i) dynamically adjusting memory
values in response to the current context input to selectively weaken
outdated or enhance relevant information, and (ii) ensuring the expected
value of the calibration matrix product remains bounded, thereby preventing
gradient vanishing or exploding. Through extensive experimentation
on POMDP benchmarks, including meta reinforcement learning, long-horizon
credit assignment, and hard memorization games, we demonstrate that
our method consistently outperforms state-of-the-art memory-based
models in terms of performance while also delivering competitive speed.
We also provide comprehensive ablation studies that offer insights
into the components and internal workings of our memory models.

## 2 Background

### 2.1 Reinforcement Learning Preliminaries

A Partially Observable Markov Decision Process (POMDP) is formally
defined as a tuple $\langle S,A,O,R,P,\gamma\rangle$, where $S$
is the set of states, $A$ is the set of actions, $O$ is the observation
space, $R:S\times A\to\mathbb{R}$ is the reward function, $P:S\times A\to\Delta(S)$
defines the state transition probabilities, and $\gamma\in[0,1)$
is the discount factor. Here, the agent does not directly observe
the true environment state $s_{t}$. Instead, it receives an observation
$o_{t}\sim O(s_{t})$, which provides partial information about the
state, often not enough for optimal decision making. Therefore, the
agent must make decisions based on its current observation $o_{t}$
and a history of previous observations, actions, and rewards $(a_{0},r_{0},o_{1},\dots,a_{t-1},r_{t-1},o_{t})$.
The history may exclude past rewards or actions.

Let us denote the input context at timestep $t$ as $x_{t}=\left(o_{t},a_{t-1},r_{t-1}\right)$
and assume that we can encode the sequence of contexts into a memory
$M_{t}=f\left(\left\{x_{i}\right\}_{i=1}^{t}\right)$. The goal
is to learn a policy $\pi(a_{t}|M_{t})$ that maximizes the expected
cumulative discounted reward:

$$ $J(\pi)=\mathbb{E}_{\pi}\left[\sum_{t=1}^{\infty}\gamma^{t}R(s_{t},a_{t})|a_{t} \sim\pi(a_{t}|M_{t}),s_{t+1}\sim P(s_{t+1}|s_{t},a_{t}),o_{t}\sim O(s_{t})\right]$ (1) $$

Thus, a memory system capable of capturing past experiences is essential
for agents to handle the partial observability of the environment
while maximizing long-term rewards.

### 2.2 Memory-Augmented Neural Networks

We focus on matrix memory $M$, and to simplify notation, we assume
it is a square matrix. Given a memory $M\in\mathbb{R}^{H\times H}$,
we usually read from the memory as:

$$ $h_{t}=M_{t}q\left(x_{t}\right)$ (2) $$

where $q$ is a query network $q:\mathbb{R}^{D}\mapsto\mathbb{R}^{H}$
to map an input context $x_{t}\in\mathbb{R}^{D}$ to a query vector.
The read value $h_{t}$ later will be used as the input for policy/value
functions. Even more important than the reading process is the memory
writing:* How can information be written into the memory to ensure
efficient and accurate memory reading?* A general formulation for
memory writing is $M_{t}=f\left(M_{t-1},x_{t}\right)$ with $f$ as
the update function that characterizes the memory models.

The simplest form of memory writing traces back to Hebbian learning
rules: $M_{t}=M_{t-1}+x_{t}\otimes x_{t}$ where $\otimes$ is outer
product Kohonen and Ruohonen (); Hebb (). Later,
researchers have proposed “fast weight” memory Marr and Thach (); Schmidhuber (); Ba et al. ():

$$ $M_{t}=M_{t-1}\lambda+\eta g\left(M_{t-1},x_{t}\right)\otimes g\left(M_{t-1},x_ {t}\right)$ (3) $$

where $g$ is a non-linear function that take the previous memory
and the current input data as the input; $\lambda$ and $\eta$ are
constant hyperparameters. On the other hand, computer-inspired memory
architectures such as Neural Turing Machine (NTM, Graves et al. ())
and Differentiable Neural Computer (DNC, Graves et al. ())
introduce more sophisticated memory writing:

$$ $M_{t}=M_{t-1}\odot\left(\boldsymbol{1}-w\left(M_{t-1},x_{t}\right)\otimes e \left(M_{t-1},x_{t}\right)\right)+w\left(M_{t-1},x_{t}\right)\otimes v\left(M_ {t-1},x_{t}\right)$ (4) $$

where $w$, $e$ and $v$ are non-linear functions that take the previous
memory and the current input data as the input to produce the writing
weight, erase and value vectors, respectively. $\odot$ is the Hadamard
(element-wise) product.

The problem with non-linear $f$ w.r.t $M$ is that the computation
must be done in recursive way, and thus being slow. Therefore, recent
memory-based models adopt simplified linear dynamics. For example,
Linear Transformer’s memory update reads Katharopoulos et al. ():

$$ $M_{t}=M_{t-1}+v\left(x_{t}\right)\otimes\frac{\phi\left(k\left(x_{t}\right) \right)}{\sum_{t}\phi\left(k\left(x_{t}\right)\right)}$ (5) $$

where $\phi$ is an activation function; $k$ and $v$ are functions
that transform the input to key and value. Recently, Beck et al. ()
have proposed matrix-based LSTM (mLSTM):

$$ $M_{t}=\mathtt{f}\left(x_{t}\right)M_{t-1}+\mathtt{i}\left(x_{t}\right)v\left(x _{t}\right)\otimes k\left(x_{t}\right)$ (6) $$

where $\mathtt{f}$ and $\mathtt{i}$ are the forget and input gates,
respectively. In another perspective inspired by neuroscience, Fast
Forgetful Memory (FFM, Morad et al. ()) employs a parallelable
memory writing, which processes a single step update as follows:

$$ $M_{t}=M_{t-1}\odot\gamma+\left(v\left(x_{t}\right)\otimes\boldsymbol{1}^{\top} \right)\odot\gamma^{t-n}$ (7) $$

where $\gamma$ is a trainable matrix, $v$ is an input transformation
function, and $n$ is the last timestep.

## 3 Methods

In this section, we begin by introducing a unified memory writing
framework that incorporates several of the memory writing approaches
discussed earlier. Next, we examine the limitations of current memory
writing approaches through an analysis of this framework. Following
this analysis, we propose specialized techniques to address these
limitations. For clarity and consistency, all matrix and vector indices
will be referenced starting from 1, rather than 0. Constant matrices
are denoted by bold numbers.

### 3.1 Hadamard Memory Framework (HMF)

We propose a general memory framework that uses the Hadamard product
as its core operation. The memory writing at time step $t$ is defined
as:

$$ $M_{t}=M_{t-1}\odot\underbrace{C_{\theta}\left(x_{t}\right)}_{C_{t}}+ \underbrace{U_{\varphi}\left(x_{t}\right)}_{U_{t}}$ (8) $$

where $C_{\theta}:\mathbb{R}^{D}\mapsto\mathbb{R}^{H\times H}$and
$U_{\varphi}:\mathbb{R}^{D}\mapsto\mathbb{R}^{H\times H}$ are parameterized
functions that map the current input $x_{t}$ to two matrices $C_{t}$
(calibration matrix) and $U_{t}$ (update matrix). Here, $\theta$
and $\varphi$ are referred to as calibration and update parameters.
Intuitively, the calibration matrix $C_{t}$ determines which parts
of the previous memory $M_{t-1}$ should be weakened and which should
be strengthened while the update matrix $U_{t}$ specifies the content
to be encoded into the memory. We specifically choose the Hadamard
product ($\odot$) as the matrix operator because it operates on each
memory element individually. We avoid using the matrix product to
prevent mixing the content of different memory cells during calibration
and update. Additionally, the matrix product is computationally slower.

There are many ways to design $C_{t}$ and $U_{t}$. Given proper
choices of $C_{t}$ and $U_{t}$, Eqs. 3-7
can be reformulated into Eq. 8. Inspired by prior “fast
weight” works, we propose a simple update matrix:

$$ $U_{\varphi}\left(x_{t}\right)=\eta_{\varphi}\left(x_{t}\right)\left[v\left(x_{ t}\right)\otimes k\left(x_{t}\right)\right]$ (9) $$

where $k$ and $v$ are trainable neural networks that transform the
input $x_{t}$ to key and value representations. $\eta_{\varphi}:\mathbb{R}^{D}\mapsto\mathbb{R}$
is a parameterized function that maps the current input $x_{t}$ to
an update gate that controls the amount of update at step $t$. For
example, if $\eta_{\varphi}\left(x_{t}\right)=0$, the memory will
not be updated with any new content, whereas $\eta_{\varphi}\left(x_{t}\right)=1$
the memory will be fully updated with the content from the $t$-th
input. We implement $\eta_{\varphi}\left(x_{t}\right)$ as a neural
network with $\mathrm{sigmoid}$ activation function.

We now direct our focus towards the design of $C_{t}$, which is the
core contribution of our work. The calibration matrix selectively
updates the memory by erasing no longer important memories and reinforcing
ongoing critical ones. In a degenerate case, if $C_{t}=\boldsymbol{1}$
for all $t$, the memory will not forget or strengthen any past information,
and will only memorize new information over time, similar to the Hebbian
Rule and Transformers. To analyze the role of the calibration matrix,
it is useful to unroll the recurrence, leading to the closed-form
equation (see proof in Appendix A.1):

$$ $M_{t}=M_{0}\prod_{i=1}^{t}C_{i}+\sum_{i=1}^{t}U_{i}\odot\prod_{j=i+1}^{t}C_{j}$ (10) $$

where $\overset{}{\prod}$ represents element-wise products. Then,
$h_{t}=M_{t}q\left(x_{t}\right)=M_{0}\prod_{i=1}^{t}C_{i}q\left(x_{t}\right)+
\sum_{i=1}^{t}U_{i}\odot\prod_{j=i+1}^{t}C_{j}q\left(x_{t}\right)$.
Calibrating the memory is important because without calibration ($C_{t}=\boldsymbol{1}$
$\forall t$), the read value becomes: $h_{t}=M_{0}q\left(x_{t}\right)+\sum_{i=1}^{t}U_{i}q\left(x_{t}\right).$
In this case, making $h_{t}$ to reflect a past context at any step
$j$ requires that $q\left(x_{t}\right)\neq 0$ and $\sum_{i\neq j}U_{i}q\left(x_{t}\right)=\sum_{i\neq j}\eta_{\varphi}\left(x_{i}
\right)v\left(x_{i}\right)\left[k\left(x_{i}\right)\cdot q\left(x_{t}\right)
\right]\approx\boldsymbol{0}$,
which can be achieved if we can find $q\left(x_{t}\right)$ such that
$k\left(x_{i}\right)\cdot q\left(x_{t}\right)\approx 0\,\,\forall i\neq j$.
Yet, this becomes hard when $T\gg H$ and $\eta_{\varphi}\left(x_{i}\right)\neq\boldsymbol{0}$
as it leads to an overdetermined system with more equations than variables.
We note that avoiding memorizing any $i$-th step with $\eta_{\varphi}\left(x_{i}\right)=0$
is suboptimal since $x_{i}$ may be required for another reading step
$t^{\prime}\neq t$.

Therefore, at a certain timestep $t$, it is critical to eliminate
no longer relevant timesteps from $M_{t}$ by calibration, i.e., $U_{i}\odot\prod_{j=i+1}^{t}C_{j}\approx\boldsymbol{0}$
for unimportant $i$ (forgetting). For example, an agent may first
encounter an important event, like seeing a color code, before doing
less important tasks, such as picking apples. When reaching the goal
requiring to identify a door matching the color code, it would be
beneficial for the agent to erase memories related to the apple-picking
task, ensuring a clean retrieval of relevant information–the color
code. Conversely, if timestep $i$ becomes relevant again at a later
timestep $t^{\prime}$, we need to recover its information, ensuring $U_{i}\odot\prod_{j=i+1}^{t^{\prime}}C_{j}\neq\boldsymbol{0}$
(strengthening), just like the agent, after identifying the door,
may return to collecting apple task.

###### Remark 1 .

In the Hadamard Memory Framework, calibration should be enabled ($C_{t}\neq\boldsymbol{1}$)
and conditioned on the input context.

Regarding computing efficiency, if $C_{t}$ and $U_{t}$ are not functions
of $M_{<t}$, we can compute the memory using Eq. 10
in parallel, ensuring fast execution. In particular, the set of products
$\left\{\prod_{j=i+1}^{t}C_{j}\right\}_{i=1}^{t}$ can be calculated
in parallel in $O\left(\log t\right)$ time. The summation can also
be done in parallel in $O\left(\log t\right)$ time. Additionally,
since all operations are element-wise, they can be executed in parallel
with respect to the memory dimensions. Consequently, the total time
complexity is $O\left(\log t\right)$. Appendix Algo. 1
illustrates an implementation supporting parallelization.

###### Remark 2 .

In the Hadamard Memory Framework, with optimal parallel implementation,
the time complexity for processing a sequence of $t$ steps is $O\left(\log t\right)$.
By contrast, without parallelization, the time complexity is $O\left(tH^{2}\right)$.

###### Remark 1 .

###### Remark 2 .

### 3.2 Challenges on Memory Calibration

The calibration matrix enables agents to either forget or enhance
past memories. However, it complicates learning due to the well-known
issues of gradient vanishing or exploding. This can be observed when
examining the policy gradient over $T$ steps, which reads:

$$ $\displaystyle\nabla_{\Theta}J\left(\pi_{\Theta}\right)$ $\displaystyle=\mathbb{E}_{s,a\sim\pi_{\Theta}}\sum_{t=0}^{T}\underbrace{\nabla _{\Theta}\log\pi_{\Theta}\left(a_{t}|M_{t}\right)}_{G_{t}\left(\Theta\right)} Adv\left(s_{t},a_{t},\gamma\right)$ (11) $$

where $\Theta$ is the set of parameters, containing $\left\{\theta,\varphi\right\}$,
$Adv$ represents the advantage function, which integrates reward
information $R\left(s_{t},a_{t}\right)$, and $G_{t}\left(\Theta\right)=\frac{\partial\log\pi_{\Theta}\left(a_{t}|M_{t}\right
)}{\partial M_{t}}\frac{\partial M_{t}}{\partial\Theta}$
captures information related to the memory. Considering main gradient
at step $t$, required to learn $\theta$ and $\varphi$:

$$ $\displaystyle\frac{\partial M_{t}}{\partial\theta}$ $\displaystyle=M_{0}\frac{\partial\prod_{i=1}^{t}C_{\theta}(x_{i})}{\partial \theta}+\sum_{i=1}^{t}U_{\varphi}\left(x_{i}\right)\odot\underbrace{\frac{ \partial\prod_{j=i+1}^{t}C_{\theta}(x_{j})}{\partial\theta}}_{G_{1}\left(i,t, \theta\right)};$ (12) $\displaystyle\frac{\partial M_{t}}{\partial\varphi}$ $\displaystyle=\sum_{i=1}^{t}\frac{\partial U_{\varphi}\left(x_{i}\right)}{ \partial\varphi}\odot\underbrace{\prod_{j=i+1}^{t}C_{\theta}(x_{j})}_{G_{2} \left(i,t,\theta\right)}$ (13) $$

We collectively refer $G_{1}\left(i,t,\theta\right)$ and $G_{2}\left(i,t,\theta\right)$
as $G_{1,2}\left(i,t,\theta\right)\in\mathbb{R}^{H\times H}$. These
terms are critical as they capture the influence of state information
at timestep $i$ on the learning parameters $\theta$ and $\varphi$.
The training challenges arise from these two terms as the number of
timesteps $t$ increases: (i) *Numerical Instability (Gradient
Exploding):* if $\exists m.k\in\left[1,H\right]\,\mathrm{s.t.}\,G_{1,2}\left(i,t,\theta\right)
\left[m,k\right]\rightarrow\infty$,
this leads to overflow, causing the gradient to become “nan”;
(ii) *Learning Difficulty (Gradient Vanishing)*: if $t\gg i_{0}$,
$\left\|G_{1,2}\left(i,t,\theta\right)\right\|\approx 0\,\,\forall i<i_{0}$,
meaning no signal from timesteps $i<i_{0}$ contributes to learning
the parameters. This is suboptimal, especially when important observations
occur early in the episode, and rewards are sparse and given at the
episode end, .i.e., $R\left(s_{t},a_{t}\right)=0\text{$\,\forall\,t\neq T$}$.

*How to design the calibration matrix $C_{\theta}\left(x\right)$
to overcome the training challenges?* A common approach is to either
fix it as hyperparameters or make it learnable parameters independent
on the input $x_{t}$ (e.g., Eqs. 3 and 7).
Unfortunately, we can demonstrate that this leads to either numerical
instability or learning difficulties as formalized in Proposition
3. In the next section, we will provide a better design
for the calibration matrix.

###### Proposition 3 .

If calibration is enabled $C_{\theta}\left(x_{t}\right)\neq\boldsymbol{1}$,
yet the calibration matrix is fixed, independent of the input $x_{t}$
($\forall t:C_{\theta}\left(x_{t}\right)=\theta\in\mathbb{R}^{H\times H}$),
numerical instability or learning difficulty will arise.

###### Proof.

See Appendix A.2
∎

###### Proposition 3 .

###### Proof.

### 3.3 Stable Hadamrad Memory (SHM)

To avoid numerical and learning problems, it is important to ensure
each element of $C_{t}$ is not always greater than 1 or smaller than
1, which ends up in their product will be bounded such that $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}C_{t}\right]\neq\left\{\boldsymbol{0},
\boldsymbol{\infty}\right\}$
as $T\rightarrow\infty$. At the same time, we want $C_{t}$ to be
a function of $x_{t}$ to enable calibration conditioned on the current
context. To this end, we propose the calibration matrix:

$$ $C_{\theta}\left(x_{t}\right)=1+\tanh\left(\theta_{t}\otimes v_{c}\left(x_{t} \right)\right)$ (14) $$

where $v_{c}:\mathbb{R}^{D}\mapsto\mathbb{R}^{H}$ is a mapping function,
and $\theta_{t}\in\mathbb{R}^{H}$ represents the main calibration
parameters. Here, we implement $v_{c}$ as a linear transformation
to map the input to memory space. Notably, the choice of $\theta_{t}$
determines the stability of the calibration. *We propose to design
$\theta_{t}$ as trainable parameters that is randomly selected from
a set of parameters $\theta$*. In particular, given $\theta\in\mathbb{R}^{L\times H}$,
we sample uniformly a random row from $\theta$, $\theta_{t}=\theta\left[l_{t}\right]$
where $l_{t}\sim\mathcal{U}\left(1,L\right)$ where $L$ is the number
of possible $\theta_{t}$. We name the design as Stable Hadamard Memory
(SHM). Given the formulation, the range of an element $z_{t}^{m,k}=C_{\theta}\left(x_{t}\right)\left[m,k\right]$
is $\left[0,2\right]$ where $m,k\in\left[1,H\right]$. We can show
one important property of this choice is that $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}C_{t}\right]\approx\boldsymbol{1}$
under certain conditions.

###### Proposition 4 .

Assume: (i) $x_{t}\sim\mathcal{N}\left(0,\Sigma_{t}\right)$, (ii)
$v_{f}$ is a linear transformation, i.e.,$v_{c}\left(x_{t}\right)=Wx_{t}$,
(iii) $\left\{z_{t}=\theta_{t}\otimes v_{c}\left(x_{t}\right)\right\}_{t=1}^{T}$
are independent across $t$. Given $C_{\theta}\left(x_{t}\right)$
defined in Eq. 14 then $\forall T\geq 0,1\leq m,k\leq H$:

$$ $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}z_{t}^{m,k}\right]=\text{$\mathbb{E}$} \left[\prod_{t=1}^{T}C_{\theta}\left(x_{t}\right)\left[m,k\right]\right]=1$ $$

###### Proof.

see Appendix A.3
∎

Assumption (i) is reasonable as $x_{t}$ can be treated as being drawn
from a Gaussian distribution, and LayerNorm can be applied to normalize
$x_{t}$ to have a mean of zero. Assumption (ii) can be realized as
we implement $v_{c}$ as a linear transformation. Assumption (iii)
is more restrictive because $\left\{x_{t}\right\}_{t=1}^{T}$ are
often dependent in RL setting, which means $\left\{z_{t}=\theta_{t}\otimes v_{c}\left(x_{t}\right)\right\}_{t=1}^{T}$
are not independent and thus, $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}z_{t}^{m,k}\right]\neq 1$.
However, by reducing the linear dependence between $z_{t}$ through
random selection of $\theta_{t}$, we can make $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}z_{t}^{m,k}\right]$
closer to $1$, and thus being bounded. Specifically, we will prove
that by using $\theta_{t}=\theta\left[l_{t}\right];l_{t}\sim\mathcal{U}\left(1,L\right)$
the Pearson Correlation Coefficient between timesteps is minimized,
as stated in the following proposition:

###### Proposition 5 .

Let $z_{t}^{m,k}=u_{t}^{m}v_{t}^{k}$ where $z_{t}^{m,k}=\left(\theta_{t}\otimes v_{c}\left(x_{t}\right)\right)\left[m,k\right]$,
$u_{t}^{m}=\theta_{t}\left[m\right]$ and $v_{t}^{k}=v_{c}\left(x_{t}\right)\left[k\right]$.
Given the Pearson correlation coefficient of two random variables
$X$ and $Y$ is defined as $\rho\left(X,Y\right)=\frac{\text{Cov}(X,Y)}{\sqrt{\text{Var}(X)}\sqrt{\text{
Var}(Y)}}$,
then $\forall v_{t}^{k},v_{t^{\prime}}^{k}$:

$$ $\left|\rho\left(u_{t}^{m}v_{t}^{k},u_{t^{\prime}}^{m}v_{t^{\prime}}^{k}\right) \right|\leq\left|\rho\left(v_{t}^{k},v_{t^{\prime}}^{k}\right)\right|$ $$

The equality holds when $u_{t}^{m}=\beta u_{t^{\prime}}^{m}$.

###### Proof.

See Appendix A.4
∎

As a result, our choice of $\theta_{t}$ outperforms other straightforward
designs for minimizing dependencies between timesteps. For instance,
a fixed $\theta_{t}$ ($\theta_{t}=\theta\in\mathbb{R}^{H}$) results
in higher dependencies because $\rho(\theta\left[m\right]v_{t}^{k},\theta\left[m\right]v_{t^{\prime}}^{k})=
\rho(v_{t}^{k},v_{t^{\prime}}^{k})\geq\rho\left(u_{t}^{m}v_{t}^{k},u_{t^{
\prime}}^{m}v_{t^{\prime}}^{k}\right)$.
In practice, even when $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}z_{t}^{m,k}\right]$
is bounded, the cumulative product can occasionally become very large
for certain episodes and timesteps, leading to overflow and disrupting
the learning process. This can be avoided by clipping the gradients.
In experiments, we implement SHM using nonparallel recursive form
(Eq. 8). The memory is then integrated into policy-gradient
RL algorithms to optimize Eq. 1, with the read value
$h_{t}$ (Eq. 2) used as input for value/policy networks.
Code will be available upon publication.

###### Proposition 4 .

###### Proof.

###### Proposition 5 .

###### Proof.

## 4 Experimental Results

We evaluate our method alongside notable memory-augmented agents in
POMDP environments. Unless stated otherwise, the context consists
of the observation and previous action. All training uses the same
hardware (single NVDIA H100 GPU), RL architecture, algorithm, training
protocol, and hyperparameters. We focus only on model-free RL algorithms.
The baselines differ only in their memory components: GRU Chung et al. (),
FWP Schlag et al. (), GPT-2 Radford et al. (), S6
Gu and Dao (), mLSTM Beck et al. (), FFM Morad et al. ()
and SHM (Ours). For tasks with a clear goal, we measure performance
using the Success Rate, defined as the ratio of episodes that reach
the goal to the total number of evaluation episodes. For tasks in
POP-Gym, we use Episode Return as the evaluation metric. We fix SHM’s
number of possible $\theta_{t}$, $L=128$, across experiments.

### 4.1 Sample-Efficient Meta Reinforcement Learning

Meta-RL targets POMDPs where rewards and environmental dynamics differ
across episodes, representing various tasks Schmidhuber (); Thrun and Pratt ().
To excel in all tasks, memory agents must learn general memory update
rules that can adapt to any environments. We enhanced the Wind and
Point Robot environments from Ni et al. () to increase
difficulty. In these environments, the observation consists of the
agent’s 2D position $p_{t}$, while the goal state $p_{g}$ is hidden.
The agent takes continuous actions $a_{t}$ by moving with 2D velocity
vector. The sparse reward is defined as $R(p_{t+1},p_{g})=\boldsymbol{1}(\|p_{t+1}-p_{g}\|_{2}\leq r)$
where $r=0.01$ is the radius. In Wind, the goal is fix $p_{g}=\left[0,1\right]$,
yet there are noises in the dynamics: $p_{t+1}=p_{t}+a_{t}+w,$
with the “wind” $w$ sampled from $U[-0.1,0.1]$ at the start
and fixed thereafter. In Point Robot, the goal varies across episodes,
sampled from $U[-10,10]$. To simulate real-world conditions where
the training tasks are limited, we create 2 modes using different
number of training and testing tasks: $\left[50,150\right]$ and $\left[10,190\right]$,
respectively. Following the modifications, these simple environments
become significantly more challenging to navigate toward the goal,
so we set the episode length to 100.

We incorporate the memory methods to the Soft Actor Critic (SAC, Haarnoja et al. ()),
using the same code base introduced in Ni et al. (). We
keep the SHM model sizes and memory capacities small, at around 2MB
for the checkpoint and 512 memory elements, which is roughly equivalent
to a GRU (see Appendix C.1). We
train all models for 500,000 environment steps, and report the learning
curves in Fig. 1. In the easy mode (50-150),
our SHM method consistently achieves the best performance, with a
near-optimal success rate, while other methods underperform by 20-50%
on average in Wind and Point Robot, respectively. In the hard mode
(10-190), SHM continues to outperform other methods by approximately
20%, showing earlier signs of learning.

Figure: Figure 1: Meta-RL: Wind and Point Robot learning curves. Mean $\pm$ std. over 5 runs.
Refer to caption: x1.png

Figure: Figure 2: Credit Assignment: Visual Match, Key-to-Door learning curves. Mean $\pm$ std. over 3 runs.
Refer to caption: x2.png

### 4.2 Long-term Credit Assignment

In this task, we select the Visual Match and Key-to-Door environments,
the most challenge tasks mentioned in Ni et al. ().
Both have observation as the local view of the agent, discrete actions
and sparse rewards dependent on the full trajectory, requiring long-term
memorization. In particular, the pixel-based Visual Match task features
an intricate reward system: in Phase 1, observing color codes yields
no rewards, while in Phase 2, picking an apple provides immediate
reward of one, relying on short-term memory. The final reward–a
bonus for reaching the door with the matching color code is set to
10. Key-to-Door also involves multiple phases: finding the key, picking
apples, and reaching the door. The terminal reward is given if the
key was acquired in the first phase and used to open the door. Both
tasks can be seen as decomposed episodic problems with noisy immediate
rewards, requiring that in the final phase, the agent remembers the
event in the first phase. We create 2 configurations using different
episode steps in the second phases: 250 and 500, respectively.

Still following Ni et al. (), we use SAC-Discrete Christodoulou ()
as the RL algorithm. We use the same set of baselines as in Sec. 4.1
and train them for 2 million environment steps. The results in Fig.
2 clearly show that, our method, SHM, significantly
outperforms all other methods in terms of success rate. Notably, SHM
is the only method that can perfectly solve both 250 and 500-step
Visual Match while the second-best performer, FFM, achieves only a
77% and 25% success rate, respectively. In Key-To-Door, our method
continues showing superior results with high success rate. By contrast,
no meaningful learning is observed from the other methods, which perform
similarly to GPT-2, as also noted by Ni et al. ().

### 4.3 POPGym Hardest Games

We evaluate SHM on the POPGym benchmark Morad et al. (), the
largest POMDP benchmark to date. Following previous studies Morad et al. (); Samsami et al. (),
we focus on the most memory-intensive tasks: Autoencode, Battleship,
Concentration and RepeatPrevious. These tasks require ultra long-term
memorization, with complexity increasing across Easy, Medium, and
Hard levels. All tasks use categorical action and observation spaces,
allowing up to 1024 steps.

For comparison, we evaluate SHM against *state-of-the-art model-free
methods*, including GRU and FFM. Other memory models, such as DNC,
Transformers, FWP, and SSMs, have been reported to perform worse.
The examined memory models are integrated into PPO Schulman et al. (),
trained for 15 million steps using the same codebase as Morad et al. ()
to ensure fairness. The models differ only in their memory, controlled
by the memory dimension hyperparameter $H$. We tune it for each baseline,
adjusting it to optimize performance, as larger $H$ values typically
improve results. The best-performing configurations are reported in
Table 1, where SHM demonstrates a
relative improvement of $\approx$10-12% over FFM and GRU on average.
Notably, the learning curves in Appendix Fig. 4
show that only SHM demonstrates signs of learning in several tasks,
including RepeatPrevious-Medium/Hard, and Autoencode-Easy/Medium.
Detailed hyperparameter setting and additional results are provided
in Appendix C.2.

In terms of running time, the average batch inference time in milliseconds
for GRU, FFM, and SHM is 1.6, 1.8, and 1.9, respectively, leading
to a total of 7, 8, and 9 hours of training per task. While SHM is
slightly slower than GRU and FFM, the difference is a reasonable trade-off
for improved memory management in partially observable environments.
Last but not least, our model’s runtime was measured using a non-parallel
implementation, while GRU benefits from hardware optimization in the
PyTorch library. SHM’s running time could be further improved with
proper parallelization.

**Table 1: PopGym: Mean return $\pm$ std. ($\times 100$) at the end of training over 3 runs. The range of return ($\times 100$) is $\left[-100,100\right]$.**
| Task | Level | GRU | FFM | SHM (Ours) |
| --- | --- | --- | --- | --- |
| Autoencode | Easy | -37.9$\pm$7.7 | -32.7$\pm$0.6 | 49.5$\pm$23.3 |
| Medium | -43.6$\pm$3.5 | -32.7$\pm$0.6 | -28.8$\pm$14.4 |  |
| Hard | -48.1$\pm$0.7 | -47.7$\pm$0.5 | -43.9$\pm$0.9 |  |
| Battleship | Easy | -41.1$\pm$1.0 | -34.0$\pm$7.1 | -12.3$\pm$2.4 |
| Medium | -39.4$\pm$0.5 | -37.1$\pm$3.1 | -16.8$\pm$0.6 |  |
| Hard | -38.5$\pm$0.5 | -38.8$\pm$0.3 | -21.2$\pm$2.3 |  |
| Concentration | Easy | -10.9$\pm$1.0 | 10.7$\pm$1.2 | -1.9$\pm$2.4 |
| Medium | -21.4$\pm$0.5 | -24.7$\pm$0.1 | -21.0$\pm$0.8 |  |
| Hard | -84.0$\pm$0.3 | -87.5$\pm$0.5 | -83.3$\pm$0.1 |  |
| RepeatPrevious | Easy | 99.9$\pm$0.0 | 98.4$\pm$0.3 | 88.9$\pm$11.1 |
| Medium | -34.7$\pm$1.7 | -24.3$\pm$0.4 | 48.2$\pm$7.2 |  |
| Hard | -41.7$\pm$1.8 | -33.9$\pm$1.0 | -19.4$\pm$9.9 |  |
| Average | All | -28.4$\pm$1.3 | -24.2$\pm$1.2 | -5.1$\pm$6.3 |

### 4.4 Model Analysis and Ablation Study

Choice of Calibration Prop. 5 suggests that
selecting random $\theta_{t}\in\mathbb{R}^{H}$ in Eq. 14
will reduce the dependencies between $C_{t}$, bringing $\prod_{t=1}^{T}C_{t}$
closer to $\boldsymbol{1}$ to avoid gradient issues. In this section,
we empirically verify that by comparing our proposed Random $\theta_{t}$
with the following designs: $C=\boldsymbol{1}$, no calibration is
used; *Random $C$*, where a random calibration matrix is sampled
from normal distribution at each timestep, having $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}z_{t}^{m,k}\right]=1$
under mild assumptions, but preventing learning meaningful calibration
rules;* Fixed $C$*, a simpler method for learning calibration,
but prone to gradient problems (Prop. 3); *Fixed $\theta_{t}$*,
where we learn fixed parameter $\theta_{t}$, which is proven to be
less effective than Random $\theta_{t}$ in reducing linear dependencies
between timesteps (Prop. 5); Neural $\theta_{t}$, where
$\theta_{t}=FFW\left(x_{t}\right)$, generated by a feedforward neural
network like mLSTM, but with no guarantee of reducing timestep dependencies.
We trained RL agents using the above designs of the calibration matrix
with $H=72$ on the Autoencode-Easy task and present the learning
curves in Fig. 3 (a, left). The results
show that our proposed Random $\theta_{t}$ outperforms the other
baselines by a substantial margin, with at least a 30% improvement
in performance. This confirms the effectiveness of our calibration
design in enhancing the final results.

Figure: Figure 3: (a) Left: Return of calibration designs over 3 runs; Right: Calibration matrix cumulative product over 100 episodes. (b) Return of memory sizes $H$ on Autoencode-Easy (left) and Battleship-Easy (right). (c) Memory ($M$, top) and calibration ($C$, bottom) matrices over timesteps in Visual Match: SHM erases memory that are no longer required and strengthens the important ones.

Vanishing Behavior In practice, exploding gradients can be
mitigated by clipping. Thus, we focus on the vanishing gradient, which
depends on the cumulative product $\mathcal{C}_{j}$ $=\prod_{t=1}^{j}C_{t}$.
Our theory suggests that Random $\theta_{t}$ should be less susceptible
to the vanishing phenomenon compared to other methods such as Fixed
$C$, Fixed $\theta_{t}$ and Neural $\theta_{t}$. To verify that,
for each episode, we compute the average value of elements in the
matrix $\mathcal{C}_{j}$ that are smaller than 1 ($\overline{\mathcal{C}_{j}}\left[<1\right]$),
as those larger than 1 are prone to exploding and are not appropriate
for averaging with the vanishing values. We plot $\overline{\mathcal{C}_{j}}\left[<1\right]$
for $j=1,2,...100$ over 100 episodes in Fig 3
(a, right).

The results are consistent with theoretical predictions: Fixed $C$
leads to rapid vanishing of the cumulative product in just 10 steps.
Neural $\theta_{t}$ also perform badly, likely due to more complex
dependencies between timesteps because $z_{t}^{m,k}$ now becomes
non-linear function of $x_{t}$, causing $\prod_{t=1}^{T}z_{t}^{m,k}$
to deviate further from 1. While Fixed $\theta_{t}$ is better than
Neural $\theta_{t}$, it still exhibits quicker vanishing compared
to our approach Random $\theta_{t}$. As expected, Random $C$ shows
little vanishing, but like setting $C=1$, it fails to leverage memory
calibration, resulting in underperformance (Fig. 3
(a, left)). Random $\theta_{t}$, although its $\overline{\mathcal{C}_{j}}$
also deviates from $1$, shows the smallest deviation among the calibration
learning approaches. Additionally, the vanishing remain manageable
after 100 timesteps, allowing gradients to still propagate effectively
and enabling the calibration parameters to be learned.

Memory Size The primary hyperparameter of our method is $H$,
determining the memory capacity. We test SHM with $H\in\left\{24,72,128,156\right\}$
on the Autoencode-Easy and Battleship-Easy tasks. Fig. 3
(b) shows that larger memory generally results in better performance.
In terms of speed, the average batch inference times (in milliseconds)
for different $H$ values are 1.7, 1.8, 1.9, and 2.1, respectively.
We choose $H=128$ for other POPGym tasks to balance performance and
speed.

Forgetting and Remembering Here, we investigate the learned
calibration strategy of SHM on Visual Match with 15,100, and 10 steps
in Phase 1, 2 and, 3, respectively. We sample several representative
$M_{t}$ and $C_{t}$ from 3 phases and visualize them in Fig. 3
(c). In Phase 1, the agent identifies the color code and stores it
in memory, possibly in two columns of $M$, marked as “important
memory”. In Phase 2, unimportant memory elements are erased where
$C_{t}\approx 0$ (e.g., those within the blue rectangle). However,
important experiences related to the Phase 1’s code are preserved
across timesteps until Phase 3 (e.g., those within the red rectangle
where $C_{t}\gtrsim$1), which is desirable.

## 5 Related works

Numerous efforts have been made to equip RL agents with memory mechanisms.
Two main research directions focus on inter-episode and intra-episode
memory approaches. While episodic control with a global, tabular memory
enhances sample efficiency by storing experiences across episodes
Blundell et al. (); Le et al. (, ), it falls short
in recalling specific events within individual episodes. Similarly,
global memory mechanisms can support exploration or optimization Badia et al. (); Le et al. (, ),
but are not designed to address the challenges of memorization within
a single episode. In contrast, to address the challenges of partially
observable Markov decision processes (POMDPs), RL agents often leverage
differentiable memory, which is designed to capture the sequence of
observations within an episode, and can be learned by policy gradient
algorithms Wayne et al. ().

Differentiable memory models can be broadly categorized into vector-based
and matrix-based approaches. Vector-based memory, like RNNs Elman (),
processes inputs sequentially and stores past inputs in their hidden
states. While RNNs are slower to train, they are efficient during
inference. Advanced variants, such as GRU and LSTM, have shown strong
performance in POMDPs, often outperforming more complex RL methods
Ni et al. (); Morad et al. (). Recently, faster alternatives
like convolutional and structured state space models (SSM) have gained
attention Bai et al. (); Gu et al. (), though their effectiveness
in RL is still under exploration. Initial attempts with models like
S4 underperformed in POMDP tasks Morad et al. (), but improved
SSM versions using S5, S6 or world models have shown promise Lu et al. (); Gu and Dao (); Samsami et al. ().
Despite these advancements, vector-based memory is limited, as compressing
history into a single vector makes it challenging to scale for high-dimensional
memory space.

Matrix-based memory, on the other hand, offers higher capacity by
storing history in a matrix but at the cost of increased complexity.
Attention-based models, such as Transformers Vaswani et al. (),
have largely replaced RNNs in SL, also delivering good results in
standard POMDPs Parisotto et al. (). However, their quadratic
memory requirements limit their use in environments with long episodes.
Empirical studies have also shown that Transformers struggle with
long-term memorization and credit assignment tasks Ni et al. ().
While classic memory-augmented neural networks (MANNs) demonstrated
good performance in well-crafted long-term settings Graves et al. (); Hung et al. (); Le et al. (),
they are slow and do not scale well in larger benchmarks like POPGym
Morad et al. (). New variants of LSTM Beck et al. (),
including those based on matrices, have not been tested in reinforcement
learning settings and lack theoretical grounding to ensure stability.

Simplified matrix memory models Katharopoulos et al. (); Schlag et al. (),
offer scalable solutions but have underperformed compared to simple
RNNs in the POPGym benchmark, highlighting the challenges of designing
efficient matrix memory for POMDPs. Recently, Fast and Forgetful Memory
(FFM, Morad et al. ()), incorporating inductive biases
from neuroscience, has demonstrated better average results than RNNs
in the benchmark. However, in the most memory-intensive environments,
the improvement remains limited. Compared to our approach, these matrix-based
memory methods lack a flexible memory calibration mechanism and do
not have robust safeguards to prevent numerical and learning issues
in extremely long episodes.

## 6 Discussion

In this paper, we introduced the Stable Hadamard Framework (SHF) and
its effective instance, the Stable Hadamard Memory (SHM), a novel
memory model designed to tackle the challenges of dynamic memory management
in partially observable environments. By utilizing the Hadamard product
for memory calibration and update, SHM provides an efficient and theoretically
grounded mechanism for selectively erasing and reinforcing memories
based on relevance. Our experiments on the POPGym and POMDP benchmarks
demonstrate that SHM significantly outperforms state-of-the-art memory-based
models, particularly in long-term memory tasks, while being fast to
execute. Although our theory suggests that SHM should be more stable
and mitigate gradient learning issues by reducing linear dependencies
between timesteps, this stability is not guaranteed to be perfect.
Further theoretical investigation is needed to validate and refine
these properties in future work.

## Appendix

### A Theoretical Results

#### A.1 Closed-form Memory Update

Given

$$ $M_{t}=M_{t-1}\odot C_{t}+U_{t}$ (15) $$

then

$$ $M_{t}=M_{0}\prod_{i=1}^{t}C_{i}+\sum_{i=1}^{t}U_{i}\odot\prod_{j=i+1}^{t}C_{j}$ (16) $$

###### Proof.

We prove by induction.

Base case: for $t=1$, the equation becomes for both Eqs.
15 and 16:

$$ $M_{1}=M_{0}\cdot C_{1}+U_{1}$ $$

Thus, the equation holds for $t=1$.

Inductive hypothesis: Assume the equation holds for $t=n$:

$$ $M_{n}=M_{0}\prod_{i=1}^{n}C_{i}+\sum_{i=1}^{n}U_{i}\odot\prod_{j=i+1}^{n}C_{j}$ $$

Inductive step: We now prove the equation holds for $t=n+1$.
From the update rule in Eq. 15:

$$ $M_{n+1}=M_{n}\odot C_{n+1}+U_{n+1}$ $$

Substitute$M_{n}$ using the inductive hypothesis:

$$ $\displaystyle M_{n+1}$ $\displaystyle=\left(M_{0}\prod_{i=1}^{n}C_{i}+\sum_{i=1}^{n}U_{i}\odot\prod_{j =i+1}^{n}C_{j}\right)\odot C_{n+1}+U_{n+1}$ $\displaystyle=M_{0}\prod_{i=1}^{n+1}C_{i}+\sum_{i=1}^{n}U_{i}\odot\prod_{j=i+1 }^{n+1}C_{j}+U_{n+1}$ $\displaystyle=M_{0}\prod_{i=1}^{n+1}C_{i}+\sum_{i=1}^{n+1}U_{i}\odot\prod_{j=i +1}^{n+1}C_{j}$ $$

This matches the form of the closed-form Eq. 16 for
$t=n+1$, completing the proof by induction.
∎

###### Proof.

#### A.2 Proposition 3

###### Definition 6 .

Critical Memory Gradients: In the Hadamard Memory Framework,
we define the critical memory gradients of memory rules as follows:

$$ $G_{1}\left(i,t,\theta\right)=\frac{\partial\prod_{j=i+1}^{t}C_{\theta}(x_{j})} {\partial\theta}=\sum_{j=i+1}^{t}\frac{\partial C_{\theta}\left(x_{j}\right)}{ \partial\theta}\odot\prod_{k=i+1,k\neq j}^{t}C_{\theta}(x_{k})$ $$

$$ $G_{2}\left(i,t,\theta\right)=\prod_{j=i+1}^{t}C_{\theta}(x_{j})$ $$

Now we can proceed to the proof for Proposition 3.

###### Proof.

In the case $C_{\theta}\left(x_{t}\right)=\theta\in\mathbb{R}^{H\times H}$,
the critical gradients read:

$$ $\displaystyle G_{1}\left(i,t,\theta\right)$ $\displaystyle=\sum_{j=i+1}^{t}\prod_{k=i+1,k\neq j}^{t}\theta$ $\displaystyle=\left(t-i\right)\theta^{t-i-1}$ $$

$$ $G_{2}\left(i,t,\theta\right)=\theta^{t-i}$ $$

If $\exists\,0\leq m,k<H\,\mathrm{s.t.}\,\left|\theta\left[m,k\right]\right|>1$,
$G_{1,2}\left(i,t,\theta\right)\left[m,k\right]\sim o\left(\theta\left[m,k
\right]^{t-i-1}\right)$,
and thus $\rightarrow\infty$, i.e., numerical problem arises. Let
$\left\|\cdot\right\|$ denote the infinity norm, we also
have:

$$ $\left\|G_{1}\left(i,t,\theta\right)\right\|=\left\|\left(t-i\right)\theta^{t-i -1}\right\|\leq\left(t-i\right)\left\|\theta\right\|^{t-i-1}$ $$

$$ $\left\|G_{2}\left(i,t,\theta\right)\right\|=\left\|\theta^{t-i}\right\|\leq \left\|\theta\right\|^{t-i}$ $$

Note that if $\forall m,k\left|\theta\left[m,k\right]\right|<1$ ,
$\left\|\theta\right\|<1$, both terms become 0 as $t-i$
increases, thus learning problem always arises. In conclusion, in
this case, to avoid both numerical and learning problems, $\forall m,k\left|\theta\left[m,k\right]\right|=1$,
which is not optimal in general.
∎

###### Definition 6 .

###### Proof.

#### A.3 Proposition 4

###### Proof.

Let $z_{t}^{m,k}=u_{t}^{m}v_{t}^{k}$ where $z_{t}^{m,k}=\left(\theta_{t}\otimes v_{c}\left(x_{t}\right)\right)\left[m,k\right]$,
$u_{t}^{m}=\theta_{t}\left[m\right]$ and $v_{t}^{k}=v_{c}\left(x_{t}\right)\left[k\right]$.
Using assumption (i) and (ii), $v_{t}^{k}$ is a Gaussian variable,
i.e., $v_{t}^{k}\sim\mathcal{N}\left(0,\mu_{t}^{k}\right)$. By definition,
$u_{t}^{m}$ is a categorical random variable that can take values
$\left\{\theta\left[m,l_{t}\right]\right\}_{l_{t}=1}^{L}$ with
equal probability $1/L$. For now, we drop the subscripts $m$, $k$
and $t$ for notation ease. The PDF of $z=uv$ can be expressed as
a mixture distribution since $u$ is categorical and can take discrete
values $\left\{u_{l}\right\}_{l=1}^{L}$. The PDF of $z$, denoted
as $f\left(z\right)$, is given by:

$$ $\displaystyle f(z)$ $\displaystyle=\sum_{l=1}^{L}P(u=u_{l})f_{u_{l}v}(z)$ $\displaystyle=\frac{1}{L}\sum_{l=1}^{L}f_{u_{l}v}(z)$ $$

where $f_{u_{l}v}(z)$ is the PDF of $u_{l}v$, and $u_{l}$ is a
constant for each $l$. Thus, $f_{u_{l}v}(z)$ is the scaled PDF of
$u$:

$$ $f_{u_{l}v}(z)=\frac{1}{|u_{l}|}f_{v}\left(\frac{z}{u_{l}}\right)$ $$

Since $v\sim\mathcal{N}\left(0,\mu\right)$, the PDF of $v$, denoted
as $f_{v}(x)$, is symmetric about 0, we have:

$$ $f_{u_{l}v}(z)=\frac{1}{|u_{l}|}f_{v}\left(\frac{z}{u_{l}}\right)=\frac{1}{|u_{ l}|}f_{v}\left(\frac{-z}{u_{l}}\right)=f_{u_{l}v}(-z).$ $$

This shows that $f_{u_{l}v}(z)$ is symmetric around 0 for each $l$.
Therefore, the PDF $f(z)$ is also symmetric:

$$ $f(z)=\frac{1}{L}\sum_{l=1}^{L}f_{u_{l}v}(z)=\frac{1}{L}\sum_{l=1}^{L}f_{u_{l}v }(-z)=f(-z)$ $$

Since $\tanh$ is an odd function and the PDF of $z^{m,k}$ is symmetric
about 0, $\mathbb{E}\left[\tanh\left(z^{m,k}\right)\right]=0$ and
thus $\mathbb{E}\left[1+\tanh\left(z^{m,k}\right)\right]=1$. Finally,
using assumption (iii), $\text{$\mathbb{E}$}\left[\prod_{t=1}^{T}z_{t}^{m,k}\right]=\prod_{t=1}^{T}
\text{$\mathbb{E}$}\left[z_{t}^{m,k}\right]=1$.
∎

###### Proof.

#### A.4 Proposition 5

###### Proof.

Without loss of generalization, we can drop the indice $m$ and $k$
for notation ease. Since $\theta_{t}=\theta\left[l_{t}\right];l_{t}\sim\mathcal{U}\left(1,L\right)$,
it is reasonable to assume that each of $\left\{u_{t},u_{t^{\prime}}\right\}$,
$\left\{u_{t},v_{t}\right\}$, $\left\{u_{t},v_{t^{\prime}}\right\}$,
$\left\{u_{t^{\prime}},v_{t}\right\}$, $\left\{u_{t^{\prime}},v_{t^{\prime}}\right\}$
are independent. In this case, let us denote $X=u_{t}v_{t}$ and $X^{\prime}=u_{t^{\prime}}v_{t^{\prime}}$,
we have

$$ $\displaystyle\text{Cov}(X,X^{\prime})$ $\displaystyle=\text{Cov}(u_{t}v_{t},u_{t^{\prime}}v_{t^{\prime}})$ $\displaystyle=\mathbb{E}[u_{t}v_{t}\cdot u_{t^{\prime}}v_{t^{\prime}}]-\mathbb {E}[u_{t}v_{t}]\cdot\mathbb{E}[u_{t^{\prime}}v_{t^{\prime}}]$ $\displaystyle=\mathbb{E}(u_{t}v_{t})\mathbb{E}(u_{t^{\prime}}v_{t^{\prime}})- \mathbb{E}(u_{t})\mathbb{E}(v_{t})\mathbb{E}(u_{t^{\prime}})\mathbb{E}(v_{t^{ \prime}})$ $\displaystyle=\mathbb{E}(u_{t})\mathbb{E}(v_{t})[\mathbb{E}(u_{t^{\prime}}v_{t ^{\prime}})-\mathbb{E}(u_{t^{\prime}})\mathbb{E}(v_{t^{\prime}})]$ $\displaystyle=\mathbb{E}[u_{t}]\mathbb{E}[u_{t^{\prime}}]\text{Cov}(v_{t},v_{t ^{\prime}})$ $$

The variances read:

$$ $\text{Var}(X)=\text{Var}(u_{t}v_{t})=\mathbb{E}[u_{t}^{2}v_{t}^{2}]-\mathbb{E} [u_{t}v_{t}]^{2}=\mathbb{E}[u_{t}^{2}]\left(\mathbb{E}[v_{t}^{2}]-\mathbb{E}[v _{t}]^{2}\right)=\mathbb{E}[u_{t}^{2}]\cdot\text{Var}(v_{t})$ $$

Similarly:

$$ $\text{Var}(X^{\prime})=\text{Var}(u_{t^{\prime}}v_{t^{\prime}})=\mathbb{E}[u_{ t^{\prime}}^{2}]\cdot\text{Var}(v_{t^{\prime}})$ $$

Given $\rho$ as the Pearson Correlation Coefficient, consider the
ratio:

$$ $\displaystyle\frac{\left|\rho(X,X^{\prime})\right|}{\left|\rho(v_{t},v_{t^{ \prime}})\right|}$ $\displaystyle=\frac{\frac{\left|\mathbb{E}[u_{t}]\mathbb{E}[u_{t^{\prime}}] \text{Cov}(v_{t},v_{t^{\prime}})\right|}{\sqrt{\mathbb{E}[u_{t}^{2}]\cdot\text {Var}(v_{t})\cdot\mathbb{E}[u_{t^{\prime}}^{2}]\cdot\text{Var}(v_{t^{\prime}}) }}}{\frac{\left|\text{Cov}(v_{t},v_{t^{\prime}})\right|}{\sqrt{\text{Var}(v_{t })\cdot\text{Var}(v_{t^{\prime}})}}}$ $\displaystyle=\frac{\left|\mathbb{E}[u_{t}]\mathbb{E}[u_{t^{\prime}}]\right|}{ \sqrt{\mathbb{E}[u_{t}^{2}]\mathbb{E}[u_{t^{\prime}}^{2}]}}$ $$

By the independence of $u_{t}$ and $u_{t^{\prime}}$ and the Cauchy-Schwarz
inequality:

$$ $\left|\mathbb{E}[u_{t}]\mathbb{E}[u_{t^{\prime}}]\right|=\left|\mathbb{E}[u_{t }u_{t^{\prime}}]\right|\leq\sqrt{\mathbb{E}[u_{t}^{2}]\mathbb{E}[u_{t^{\prime} }^{2}]},$ $$

which implies

$$ $\frac{\left|\rho(u_{t}v_{t},u_{t^{\prime}}v_{t^{\prime}})\right|}{\left|\rho(v _{t},v_{t^{\prime}})\right|}\leq 1\Longleftrightarrow\left|\rho\left(u_{t}v_{t },u_{t^{\prime}}v_{t^{\prime}}\right)\right|\leq\left|\rho\left(v_{t},v_{t^{ \prime}}\right)\right|$ $$

The equality holds when $u_{t}=\beta u_{t^{\prime}}$.
∎

###### Proof.

### B Details on Methodology

In this section, we describe RL frameworks used across experiments,
which are adopted exactly from the provided benchmark. Table 2
summarizes the main configurations. Further details can be found in
the benchmark papers Ni et al. [], Morad et al. [].

**Table 2: Network architecture shared across memory baselines.**
| Task | Input type | Policy/Value networks | RL algorithm | Batch size |
| --- | --- | --- | --- | --- |
| Meta-RL | Vector | 3-layer FFW + 1 Memory | SAC | 32 |
| layer $(128,128,128,H)$ |  |  |  |  |
| Credit | Image | 2-layer CNN + 2-layer FFW | SAC-D | 32 |
| Assignment | + 1 Memory layer $(128,128,H)$ |  |  |  |
| POPGym | Vector | 3-layer FFW + 1 Memory | PPO | 65,536 |
| layer $(128,64,H,64)$ |  |  |  |  |

Figure: Algorithm 1 Theoretical Parallel Hadamard Memory Framework.

### C Details of Experiments

We adopt public benchmark repositories to conduct our experiments.
The detail is given in Table 3

**Table 3: Benchmark repositories used in our paper.**
| Task | URL | License |
| --- | --- | --- |
| Meta-RL | [https://github.com/twni2016/pomdp-baselines](https://github.com/twni2016/pomdp-baselines) | MIT |
| Credit Assignment | [https://github.com/twni2016/Memory-RL](https://github.com/twni2016/Memory-RL) | MIT |
| POP-Gym | [https://github.com/proroklab/popgym](https://github.com/proroklab/popgym) | MIT |

#### C.1 Sample-Efficiency in Meta Reinforcement Learning and Long-term Credit
Assignment Details

We report the choice of memory hyperparameter $H$ in Table. 4.
Due to the lack of efficient parallel processing in the codebase from
Ni et al. [], running experiments with larger memory sizes
is prohibitively slow. As a result, we were only able to test models
with 512 memory elements, limiting the potential performance of our
SHM. This constraint contrasts with the POPGym benchmark with better
parallel processing, where our method scales comfortably with larger
memory sizes, as demonstrated later in C.2.

**Table 4: Memory dimension for Meta-RL and Credit Assignment tasks. GRU and S6 use vector memory of size $H$. GPT-2 does not have a fixed size memory and attend to all previous $T$ timesteps in the episode. $H=512$ is the dimension of Transformer’s hidden layer. FFM’s memory shape is $2\times m\times c$ where $c=4$, $H=m=128$. FWP and SHM’s memory shape is $H\times H$.**
| Model | GRU | FWP | GPT-2^∗ | S6 | mLSTM | FFM | SHM (Ours) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| $H$ | 512 | 24 | 512 | 512 | 24 | 128 | 24 |
| Memory elements | 512 | 512 | $512\times T$ | 512 | 512 | 512 | 512 |

#### C.2 POPGym Hardest Games Details

In this experiment, we tuned the memory model size to ensure optimal
performance. Specifically, for GRU, we tested hidden sizes of 256,
512, 1024, and 2048 on the Autoencode-Easy task. Increasing GRU’s
hidden size did not lead to performance gains but resulted in a significant
rise in computation cost. For instance, at $H=2048$, the average
batch inference time was 30 milliseconds, compared to 1.6 milliseconds
at $H=256$. Thus, we set GRU’s hidden size to 256, as recommended
by the POPGym documentation. For FFM, we set the context size to $c=4$,
as the authors suggest that a larger context size degrades performance.
We tuned the trace size $m\in\left\{32,128,256,512\right\}$ on
the Autoencode-Easy task and found that $m=128$ provided slightly
better results, so we used this value for all experiments. For our
method, we tested $H\in\left\{24,72,128,156\right\}$ on the same
task and observed that larger values of $H$ led to better performance,
as shown in the ablation study in Sec. 4.4.
However, to balance runtime, memory cost, and performance, we set
$H=128$ for all POPGym experiments. The learning curves are given
in Fig. 4.

Figure: Figure 4: POPGym learning curves: Mean $\pm$ std. over 3 runs.
Refer to caption: x3.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: * Typo in the title of Section 3.3 - Stable Hadamrad Memory.

* How were the baselines tuned for each of the experiments? I could only find partial detail in the appendix. What efforts were made to normalize computational and memory cost between baselines and SHM?
soundness: 2
strengths: The paper is well presented and the method is explained in detail. It is evaluated against a range of baselines and on multiple benchmarks.
summary: This work presents a memory-based architecture for partially observable RL, named Stable Hadamard Memory (SHM). SHM is evaluated on a range of POMDP environments and POPGym against recurrent and attention-based baselines, as well as some ablations and analysis.
weaknesses: 1. The novelty of the method is not directly apparent from reading the paper. Stable Hadamard Memory is a confusing description for the method, as many existing RNNs, such as GRUs, make use of a Hadamard product as a core component. Furthermore, the Hadamard product is not the defining feature of the method, any more than addition. Instead, it would seem that the design of the calibration matrix, with random selection/dropout of update parameters, is the primary contribution of the approach. Based on this, more focus on this component in the paper and method naming would benefit the reader.

2. The method's update rule is highly similar to linear attention, which is cited but not discussed in the work. A discussion of how this method relates to linear attention and empirical evaluation against it would be compelling.

3. Due to the random selection of update parameters in SHM, the memory usage of SHM should be higher than methods without parameter dropout. Discussing this and adding experiments comparing SHM to baselines with the same memory usage would be useful.

4. The experiments have very few repeats (3 seeds) and it is unclear how some baselines were tuned.

5. A paper from RLC 2024 exploring memory aggregation in Meta-RL is not cited [1].

[1] Jacob Beck, Matthew Jackson, Risto Vuorio, Zheng Xiong, and Shimon Whiteson. Splagger: Split aggregation for meta-reinforcement learning. Reinforcement Learning Conference, 2024.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: Please see weaknesses.
soundness: 3
strengths: * As said in my summary, the results show improvements on a broad set of tasks and appear to be extremely significant.

* The method is mathematically well motivated, justified well and theoretically guarantees are given. I especially appreciate the stability guarantees.  It appears to be sound to use this for updating agent memories.

* The credit assignment benefits are justified well and empirically shown.

* The work is very well motivated and many of the theoretical guarantees are strong and supported well.

Overall this work is solid and presents and interesting way of updating memory.  I currently support acceptance of this work.  I do not have high confidence in my review as this work is a bit orthogonal to the previous works I have done on partially observably environments.
summary: This work addresses memory in partially observable RL environments. The authors state that current methods fail to capture relevant past information, making it difficult to retain important memories and know when to dynamically update them in response to changes in the environment. This issue is particularly significant for continual learning.

The authors introduce a Hadamard memory framework, suggesting that it performs memory calibration and linearly adjusts memory elements by maintaining a calibration matrix used to update these elements via an update matrix. The Hadamard product allows these elements to be learned using neural networks, and this method is well-supported in the literature. The memory matrix is then formulated using the standard policy gradient method, and the authors also demonstrate stability guarantees.

This approach is based on the Hadamard product applied to the memory matrix, enabling adaptive memory rules to be learned.

The method is evaluated on various Meta-RL tasks and challenging credit assignment tasks, with results showing dramatic improvements. The authors propose that this dynamic memory writing method facilitates more efficient credit assignment. They also evaluate the method on a set of partially observable gym environments, visually demonstrating the memory and calibration matrices.
weaknesses: * The paper is overly complex at times going into too many mathematical details.  I would like a very clear description of the benefits of using this update rule and why it empirically works well.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 2
questions: - Why is $1 + \tanh(x)$ used instead of $\textrm{sigmoid}(x)$?
- How come the recurrent state is real-valued compared to methods like FFM? Did the authors try using complex-valued states to track phase information across sequences?
- Line 233: What is $m.k$?
soundness: 3
strengths: - The paper is nicely structured
- The authors provide a thorough literature review
- The authors' method is very straightforward to implement
- The derivation of the linear matrix C is much shorter and easier to read than the derivation of the A and B matrices from the State Space Models paper (and yet, the authors' method performs better)
- Rather than designing linear operations to keep inputs around in the recurrent state as long as possible (State Space Models), the authors take relatively novel approach to design a linear operation that has well-behaved gradients
- Randomly sampling $\theta$ to break the dependence between inputs is a novel and interesting idea
- The authors run numerous experiments, across multiple benchmarks and RL algorithms
summary: The authors propose a relatively simple recurrent model for reinforcement learning under partial observability. Their model is parallelizable in time due to a linear recurrent state update. They demonstrate that their method outperforms prior work across various benchmarks, both with PPO SAC algorithms.

All in all, this is a good paper tackling an important problem with a novel approach and thorough experimentation. In my opinion, the biggest weakness is that the authors have not provided their code.
weaknesses: - The authors do not provide code
- The writing is a bit awkward in places, hampering my understanding -- I think the paper would benefit from having another proofread
    - The initialization of $C$ is quite important (lines 259-264), but it took me a few reads to understand the process -- perhaps the authors can consider writing something like $\theta = [\theta_1, \theta_2, \dots, \theta_{L}]; \quad \theta_t \sim \mathcal{U}(\theta)$
    - Proposition 3 is poorly worded, and should be rewritten
    - Section 3.3 is misspelled (STABLE HADAMRAD MEMORY (SHM))

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Hopefully, author interpret HaM in aspect of  "scaling law".

① Performance vs. Memory Size: How does the performance of the Hadamard Memory model scale with increases in memory size? Is there an optimal memory size at which the model's performance peaks?

② Impact of Task Complexity: As task complexity increases, how do the computational requirements and performance of the model change? Can the authors provide metrics to measure the impact of task complexity on model performance?

③ Scalability with Episode Length: Does the Hadamard Memory model maintain stable and efficient updates in longer episodes? Have the authors considered the model's stability and efficiency over extended periods?
soundness: 4
strengths: This paper's got some fresh ideas that really matter in the real world. First off, they've come up with this nifty thing called the Stable Hadamard Memory (SHM) model. It's like a brain upgrade for AI in situations where it can't see everything. The unique twist is this memory update method that uses something called the Hadamard product. It's all about keeping the good memories and dumping the useless ones, which makes the AI's memory super stable and efficient.

But that's not all. The paper really shows off the SHM's chops by putting it through its paces in a bunch of tests.  It outperforms the current bigwigs in memory models by a mile. This model's got the chops to handle those tough, long-term memory tasks that other models can't touch, especially when it comes to memory-intensive tasks in reinforcement learning.

So, in a nutshell, the research in this paper could really move the needle in the field of reinforcement learning.
summary: This paper introduce"SHM"model, spcefically solving memory management problem in RL. SHM use a Hadamard product-based memory update mechanism to dynamically retain important information and discard irrelevant data, aiming to enhance memory stability and efficiency. The authors propose SHM as a solution to the limitations of existing memory models, such as numerical instability and inefficient memory retention, particularly in complex and long-horizon tasks. Experimental results demonstrate SHM’s superior performance compared to state-of-the-art memory models across multiple benchmarks, including meta-reinforcement learning and the POPGym suite, indicating its potential as a robust memory mechanism for RL applications.
weaknesses: In reinforcement learning, especially when dealing with "partially observable" environments, effective memory management is really crucial. As task complexity increases, memory demand also skyrockets, and challenges grow alongside it. So, discussing how models "hold up" as tasks become more complex and expansive is pretty important.

Meta-review: This paper proposes a memory-based architecture for solving POMDPs, centered on a "memory matrix" that can be updated via a Hadamard product, which addresses numerical stability issues and efficiency issues with prior architectures. Empirical results on meta-RL tasks and credit-assignment tasks show significant improvements.

The paper is simple and well motivated, and provides strong theoretical guarantees coupled with strong empirical performance. The reviewers appreciated the paper writing, finding it "nicely structured" and noting that the derivation was "easier to read" than the SSM paper.

The main reviewer suggestions were to discuss additional prior work (other methods that use a Hadamard product; linear attention; Splagger). In terms of experiments, reviewers encouraged the authors to release code and include additional random seeds.

Taken together, this paper contributes strong theoretical and empirical results which the reviewers unanimously voted to accept. I therefore recommend that the paper be accepted.

Decision: Accept (Poster)