SlowGuess/NeurIPS_2023_OCR huggingface dataset:
If  $\xi \in \mathbb{R}_{+},\pmb{x}_j\neq 0$ , then the recurrent state will explode over long sequences:  $\lim_{n\to \infty}e^{\xi (n - j)}\pmb{x}_j = \infty$ , so the real component of  $\xi$  should be negative. Using different decay rates  $\xi_j,\xi_k\in \mathbb{R}_-$ , we can deduce the temporal ordering between terms  $x_{j}$  and  $x_{k}$ . Unfortunately, this requires both  $x_{j},x_{k}$  eventually decay to zero – what if there are important memories we do not want to forget, while simultaneously retaining their ordering? If  $\xi$  is imaginary, then we can determine the relative time between  $\pmb{x}_j,\pmb{x}_k$  as the terms oscillate indefinitely without decaying. Thus, we use a complex  $\xi$  with a negative real component, combining both forgetting and long-term temporal context. In the following paragraphs, we show that with a complex  $\xi$ , the FFM cell becomes a universal approximator of convolution.
Universal Approximation of Convolution Here, we show that FFM can approximate any temporal convolutional. Let us look at the  $z$  term from Equation 17, given the input  $\tilde{x}$ , with a slight change
Model	Training	Inference
Parallel Time	Space	Time	Space
RNN	O(n)	O(n)	O(1)	O(1)
Transformer	O(log n)	O(n2)	O(n)	O(n2)
FFM (ours)	O(log n)	O(n)	O(1)	O(1)
Table 1: The time and space complexity of memory models for a sequence of length  $n$  (training), or computing a single output recurrently during a rollout (inference).
in notation to avoid the overloaded subscript notation:  $\tilde{\pmb{x}} (\tau) = \tilde{\pmb{x}}_{\tau}$ ,  $\pmb {z}(n) = \pmb{z}_n$ . For a sequence 1 to  $n$ ,  $\pmb {z}(n)$ , the precursor to the Markov state  $\pmb{y}_n$ , can be written as:
$$
\boldsymbol {z} (n) = \boldsymbol {b} + \boldsymbol {A} \sum_ {\tau = 1} ^ {n} \tilde {\boldsymbol {x}} (\tau) \odot \exp (- \tau \boldsymbol {\alpha}) \exp (- \tau i \boldsymbol {\omega} ^ {\top}). \tag {22}
$$
where  $A, b$  is the weight and bias from linear layer  $\ell_3$ , indexed by subscript. Looking at a single input dimension  $k$  of  $\tilde{\boldsymbol{x}}$ , we have
$$
\begin{array}{l} \boldsymbol {z} (n) = \boldsymbol {b} + \sum_ {j = k c m} ^ {(k + 1) c m} \boldsymbol {A} _ {j} \sum_ {\tau = 1} ^ {n} \tilde {x} _ {k} (\tau) \exp \left(- \tau \left(i \boldsymbol {\omega} + \alpha_ {k}\right)\right) (23) \\ = \boldsymbol {b} + \sum_ {\tau = 1} ^ {n} \tilde {x} _ {k} (\tau) \sum_ {j = k c m} ^ {(k + 1) c m} \boldsymbol {A} _ {j} \exp (- \tau (i \boldsymbol {\omega} + \alpha_ {k})) (24) \\ \end{array}
$$
Equation 24 is a temporal convolution of  $\tilde{x}_k(t)$  using a Fourier Series filter with  $c$  terms  $(\omega \in \mathbb{C}^c)$ , with an additional learnable "filter extent" term  $\alpha_{k}$ . The Fourier Series is a universal function approximator, so FFM can approximate any convolutional filter over the signal  $\tilde{\pmb{x}} (\tau)$ . Appendix A further shows how this is related Laplace transform. Unlike discrete convolution, we do not need to explicitly store prior inputs or engage in zero padding, resulting in better space efficiency. Furthermore, the filter extent  $\alpha_{k}$  is learned and dynamic - it can expand for sequences with long-term temporal dependencies and shrink for tasks with short-term dependencies. Discrete temporal convolution from methods like Bai et al. [2018] use a fixed-size user-defined filter extent.
Interpretability Unlike other memory models, FFM is interpretable. Each dimension in  $S$  has a known decay rate and contextual period. We can determine trace durability (how long a trace lasts)  $t_{\alpha}$  and the maximum contextual period  $t_{\omega}$  of each dimension in  $S$  via
$$
\boldsymbol {t} _ {\alpha} = \frac {\log (\beta)}{\boldsymbol {\alpha} \mathbf {1} _ {c} ^ {\top}}, \boldsymbol {t} _ {\omega} = \frac {2 \pi}{\mathbf {1} _ {m} \boldsymbol {\omega}} \tag {25}
$$
where  $\beta$  determines the strength at which a trace is considered forgotten. For example,  $\beta = 0.01$  would correspond to the time when the trace  $x_{k}$  contributes  $1\%$  of its original value to  $S_{n}$ . FFM can measure the relative time between inputs up to a modulo of  $t_{\omega}$ .
We evaluate FFM on the two largest POMDP benchmarks currently available: POPGym [Morad et al., 2023] and the POMDP tasks from [Ni et al., 2022], which we henceforth refer to as POMDP-Baselines. For POPGym, we train a shared-memory actor-critic model using recurrent using Proximal Policy Optimization [Schulman et al., 2017]. For POMDP-Baselines, we train separate memory models for the actor and critic using recurrent Soft Actor Critic (SAC) [Haarnoja et al., 2018] and recurrent Twin Delayed DDPG (TD3) [Fujimoto et al., 2018]. We replicate the experiments from the POPGym and POMDP-Baselines papers as-is without changing any hyperparameters. We use a single FFM configuration across all experiments, except for varying the hidden and recurrent sizes to match the RNNs, and initialize  $\alpha, \omega$  based on the max sequence length. See Appendix D for further details.
Table 2 lists memory baselines, Table 3 ablates each component of FFM, Figure 3 contains a summary of training statistics for all models (wall-clock train time, mean reward, etc), and Figure 4
Name	Type	Paper/Description	Name	Type	Paper/Description
MLP	Basic	2 layer perceptron	LMU	RNN	Voelker et al. [2019]
PosMLP	Basic	MLP w/ positional encoding	IndRNN	RNN	Li et al. [2018]
Stack	Basic	Mnih et al. [2015]	Elman	RNN	Elman [1990]
TCN	Conv	Bai et al. [2018]	GRU	RNN	Chung et al. [2014]
FWP	XFormer	Schlag et al. [2021]	LSTM	RNN	Hochreiter and Schmid-huber [1997]
FART	XFormer	Katharopoulos et al. [2020]
S4D	Hybrid	Gu et al. [2021]	DNC	MANN	Graves et al. [2016]
Figure 3: Training statistics computed over ten trials and one epoch of PPO, with FFM in orange. FFM trains nearly two orders of magnitude faster on the GPU than the fastest RNNs (note the log-scale x axis). Even on the CPU, FFM trains roughly one order of magnitude faster than RNNs on the GPU. FFM has sub-milisecond CPU inference latency, making it useful for on-policy algorithms. FFM memory usage is on-par with RNNs, providing scalability to long sequences. Despite efficiency improvements, FFM still attains greater episodic reward than other models on POPGym.
contains a detailed comparison of FFM against the best performing RNN and transformer on POPGym. Figure 5 evaluates FFM on the POMDP-Baselines tasks and Figure 6 investigates FFM interpretability. See Appendix E and Appendix F for more granular plots. Error bars in all cases denote the  $95\%$  bootstrapped confidence interval.
Practicality and Robustness: FFM demonstrates exceptional performance on the POPGym and POMDP-Baseline benchmarks, achieving the highest mean reward as depicted in Figure 3 and Figure 5 without any changes to the default hyperparameters. FFM also executes a PPO training epoch nearly two orders of magnitude faster than a GRU (Figure 3, Table 1). Without forgetting, FFM underperforms the GRU, showing that forgetting is the most important inductive prior (Table 3). We use a single FFM configuration for each benchmark, demonstrating it is reasonably insensitive to hyperparameters. All in all, FFM outperforms all other models on average, across three RL algorithms and 52 tasks. Surprisingly, there are few occurrences where FFM is noticeably worse than others, suggesting it is robust and a good general-purpose model.
Explainability and Prior Knowledge: Figure 6 demonstrates that FFM learns suitable and interpretable decay rates and context lengths. We observe separate memory modes for the actor and
Table 2: The memory models used in the POPGym experiments, see Morad et al. [2023] for further details. Conv stands for convolution, XFormer stands for transformer, and MANN stands for memory augmented neural network.
FFM	FFM-NI	FFM-NO	FFM-NC	FFM-FC	FFM-ND	FFM-FD	FFM-HP
μ	0.399	0.392	0.395	0.382	0.383	0.333	0.390	0.395
σ	0.005	0.001	0.006	0.005	0.005	0.005	0.006	0.010
Table 3: Ablating FFM components over all POPGym tasks. FFM-NI removes input gating, FFM-NO removes the output gating, FFM-NC does not use temporal context, FFM-FC uses fixed (nonlearned) context, FFM-ND does not use decay, and FFM-FD uses fixed (non-learned) decay. FFM-HP uses the Hadamard product instead of the outer product to compute  $\gamma$  in Equation 12. FFM-HP requires additional parameters but ensures independence between rows and columns of the state. FFM-ND (no decay) validates our assumption that forgetting is an important inductive bias.
Figure 4: FFM compared to the GRU, the best performing model on POPGym. See Appendix F for comparisons across all other POPGym models, such as linear transformers. Error bars denote the bootstrapped  $95\%$  confidence interval over five trials. FFM noticeably outperforms state of the art on seventeen tasks, while only doing noticeably worse on four tasks.
Figure 5: Comparing FFM with tuned RNNs on the continuous control tasks from POMDP-Baselines. Ni et al. [2022] selects the best performing RNN (either GRU or LSTM) and tunes hyperparameters on a per-task basis. Error bars denote the bootstrapped  $95\%$  confidence interval over five random seeds. The task suffix V or P denotes whether the observation is velocity-only (masked position) or position-only (masked velocity) respectively. Untuned FFM outperforms tuned RNNs on all but one task for SAC. For TD3, FFM meets or exceeds RNN performance on all but one task.
Figure 6: A case study on FFM explainability. We investigate the task RepeatPreviousMedium, where the episode length is 104 timesteps and the agent must output the observation from 32 timesteps ago. (Left two) We visualize  $\pmb{\gamma}$  using trace durability  $t_{\alpha}, \beta = 0.1$  and context period  $t_{\omega}$ . Both  $t_{\alpha}, t_{\omega}$  demonstrate modes for the actor and critic, which we highlight in the figures. The critic requires the episode's current timestep to accurately predict the discounted return, while the actor just requires 32 timesteps of memory. (Left) The  $t_{\omega}$  critic mode converges to the maximum episode length of 104. The actor mode converges to 2x the necessary period (perhaps because  $\cos \left(\frac{2\pi}{t_{\omega}}\right)$ , the real component of  $\exp - ti\omega$ , is monotonic on the half-period  $0.5t_{\omega}$ ). (Middle) The  $t_{\alpha}$  terms converge to 32 for the actor and a large value for the critic. (Right) We initialize  $\alpha, \omega$  to encapsulate ideal actor and critic modes, such that  $t_{\alpha} \in [32, 104]$ ,  $t_{\omega} \in [32, 104]$ . We find FFM with informed initialization outperforms FFM with the original initialization of  $t_{\alpha} \in [1, 1024]$ ,  $t_{\omega} \in [1, 1024]$ .
critic, and we can initialize  $\alpha, \omega$  using prior knowledge to improve returns (Figure 6). This sort of prior knowledge injection is not possible in RNNs, and could be useful for determining the number of "burn-in" steps for methods that break episodes into fixed-length sequences, like R2D2 [Kapturowski et al., 2019] or MEME [Kapturowski et al., 2023].
Scaling Up: We found that large context sizes  $c$  resulted in decreased performance. We hypothesize that a large number of periodic functions results in a difficult to optimize loss landscape with many local extrema. Interestingly, transformers also employ sinusoidal encodings, which might explain the difficulty of training them in model-free RL. We tried multiple configurations of FFM cells in series, which helped in some tasks but often learned much more slowly. In theory, serial FFM cells could learn a temporal feature hierarchy similar to the spatial hierarchy in image CNNs. In many cases, serial FFM models did not appear fully converged, so it is possible training for longer could solve this issue.
Additional Experiments: Hyperparameter tuning would almost guarantee better performance, but would also result in a biased comparison to other models. We evaluated FFM with on-policy and off-policy algorithms but did not experiment with offline or model-based RL algorithms. In theory, FFM can run in continuous time or irregularly-spaced intervals simply by letting  $t$  be continuously or irregular, but in RL we often work with discrete timesteps at regular intervals so we due not pursue this further.
Numerical Precision: FFM experiences a loss of numerical precision caused by the repeated multiplication of exponentials, resulting in very large or small numbers. Care must be taken to prevent overflow, such as upper-bounding  $\alpha$ . Breaking a sequence into multiple forward passes while propagating recurrent state (setting  $n - k$  less than the sequence length) fixes this issue, but also reduces the training time efficiency benefits of FFM. We found FFM performed poorly using single precision floats, and recommend using double precision floats during multiplication by  $\gamma$ . We tested a maximum sequence length of 1024 per forward pass, although we could go higher by decreasing the maximum decay rate. With quadruple precision floats, we could process sequences of roughly 350,000 timesteps in a single forward pass with the max decay rate we used. Unfortunately, Pytorch does not currently support quads.
The inductive priors underpinning FFM are key to its success, constraining the optimization space and providing parallelism-enabling structure. Unlike many memory models, FFM is interpretable and even provides tuning opportunities based on prior knowledge about individual tasks. FFM provides a low-engineering cost "plug and play" upgrade to existing recurrent RL algorithms, improving model efficiency and reward in partially observable model-free RL with a one-line code change.
Steven Morad and Stephan Liwicki gratefully acknowledge the support of Toshiba Europe Ltd. Ryan Kortvelesy and Amanda Prorok were supported in part by ARL DCIST CRA W911NF-17-2-0181.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, and others. Human-level control through deep reinforcement learning. nature, 518(7540):529-533, 2015. Publisher: Nature Publishing Group.
Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2):179-211, 1990. Publisher: Wiley Online Library.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska-Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, and others. Hybrid computing using a neural network with dynamic external memory. Nature, 538(7626):471-476, 2016. Publisher: Nature Publishing Group.
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2019.
A Relationship to the Laplace Transform
We often think of RL in terms of discrete timesteps, but the sum in Equation 11 could be replaced with an integral for continuous-time domains, where  $\tilde{\pmb{x}}$  is a function of time:
$$
\boldsymbol {S} (n) = \int_ {0} ^ {n} \tilde {\boldsymbol {x}} (\tau) \odot \exp (- \tau \boldsymbol {\alpha}) \exp (- \tau i \boldsymbol {\omega}) ^ {\top} d \tau . \tag {26}
$$
If we look at just a single dimension of  $\tilde{x}$ , we have
$$
\boldsymbol {S} (n) = \int_ {0} ^ {n} \tilde {x} (\tau) \exp (- \tau (\alpha + i \omega)) d \tau , \tag {27}
$$
which is equivalent to the Laplace Transform  $L$  of a function  $\tilde{x}(\tau)$ . One way to interpret the connection between our aggregator and the Laplace transform is that the aggregator transforms the memory task into a pole-finding task in the S-plane, a well-known problem in control theory. When we train FFM for an RL task, we are attempting to find the poles and zeros that minimize the actor/critic loss for a batch of sequences.
B Weight Initialization
We find that our model is sensitive to initialization of the  $\alpha, \omega$  values. We compute the upper limit for  $\alpha$  such that memory will retain  $\beta = 0.01 = 1\%$  of its value after some elapsed number of timesteps we call  $t_e$ . This can be computed mathematically via
$$
\frac {\log \beta}{t _ {e}}. \tag {28}
$$
The lower limit is set based on the maximum value a double precision float can represent, minus a small  $\epsilon$
$$
\frac {\log 1 . 7 9 \times 1 0 ^ {3 0 8}}{t _ {e}} - \epsilon . \tag {29}
$$
Note that if needed, we can choose a smaller lower bound, but the input must be chunked into sequences of length  $t_e$  and run in minibatches. Since we can compute minibatches recurrently, the gradient spans across all minibatches rather than being truncated like a quadratic transformer. Nonetheless, we did not need to do this in any of our experiments. We set  $\alpha$  to be linearly spaced between the computed limits
$$
\boldsymbol {\alpha} = \left[ \alpha_ {1}, \dots \alpha_ {j} \dots \alpha_ {n} \right] ^ {\top} \tag {30}
$$
$$
\alpha_ {j} = \frac {j}{m} \frac {\log \beta}{t _ {e}} + \left(1 - \frac {j}{m}\right) \left(\frac {\log 1 . 7 9 \times 1 0 ^ {3 0 8}}{t _ {e}} - \epsilon\right) \tag {31}
$$
We initialize the  $\omega$  terms using a similar approach, with linearly spaced denominators between one and  $t_e$
$$
\omega = 2 \pi / \left[ \omega_ {1}, \omega_ {2}, \dots , \omega_ {n} \right] ^ {\top} \tag {32}
$$
$$
\omega_ {j} = \frac {j}{c} + \left(1 - \frac {j}{c}\right) t _ {e} \tag {33}
$$
All other parameters are initialized using the default Pytorch initialization.
C Computing States in Parallel with Linear Space Complexity
Here, we show how the recurrent formulation can be computed in parallel. A naive implementation would use  $O(n^{2})$  space, but using a trick, we can accomplish this in  $O(n)$  space.
Assume we have already computed hidden state  $S_{k-1}$  for some sequence. Our job is now to compute the next  $n-k+1$  recurrent states  $S_k, S_{k+1}, \ldots, S_n$  in parallel. We can rewrite Equation 11 in matrix
$$
\boldsymbol {S} _ {k: n} = \left[ \begin{array}{c} \boldsymbol {S} _ {k} \\ \boldsymbol {S} _ {k + 1} \\ \boldsymbol {S} _ {k + 2} \\ \vdots \\ \boldsymbol {S} _ {n} \end{array} \right] = \left[ \begin{array}{c} \boldsymbol {\gamma} ^ {0} \odot (\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}) + \boldsymbol {\gamma} ^ {1} \odot \boldsymbol {S} _ {k - 1} \\ \boldsymbol {\gamma} ^ {0} \odot (\boldsymbol {x} _ {k + 1} \mathbf {1} _ {c} ^ {\top}) + \boldsymbol {\gamma} ^ {1} \odot (\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}) + \boldsymbol {\gamma} ^ {2} \odot \boldsymbol {S} _ {k - 1} \\ \boldsymbol {\gamma} ^ {0} \odot (\boldsymbol {x} _ {k + 2} \mathbf {1} _ {c} ^ {\top}) + \boldsymbol {\gamma} ^ {1} \odot (\boldsymbol {x} _ {k + 1} \mathbf {1} _ {c} ^ {\top}) + \boldsymbol {\gamma} ^ {2} \odot (\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}) + \boldsymbol {\gamma} ^ {3} \odot \boldsymbol {S} _ {k - 1} \\ \vdots \\ \left(\sum_ {j = k} ^ {n} \boldsymbol {\gamma} ^ {j} \odot (\boldsymbol {x} _ {k + j} \mathbf {1} _ {c} ^ {\top})\right) + \boldsymbol {\gamma} ^ {n + 1} \odot \boldsymbol {S} _ {k - 1} \end{array} \right] \tag {34}
$$
We can write the closed form for the  $p$ th row of the matrix, where  $0 \geq p \geq n - k$
$$
\boldsymbol {S} _ {k + p} = \left(\sum_ {j = 0} ^ {p} \boldsymbol {\gamma} ^ {p - j} \odot \left(\boldsymbol {x} _ {k + j} \mathbf {1} _ {c} ^ {\top}\right)\right) + \boldsymbol {\gamma} ^ {p + 1} \odot \boldsymbol {S} _ {k - 1} \tag {35}
$$
Unfortunately, it appears we will need to materialize  $\frac{(n - k + 1)^2}{2}$  terms:
$$
\boldsymbol {\gamma} ^ {0} \odot \left(\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}\right) \tag {36}
$$
$$
\boldsymbol {\gamma} ^ {1} \odot \left(\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}\right), \boldsymbol {\gamma} ^ {0} \odot \left(\boldsymbol {x} _ {k + 1} \mathbf {1} _ {c} ^ {\top}\right) \tag {37}
$$
$$
\vdots \tag {38}
$$
$$
\boldsymbol {\gamma} ^ {n - k} \odot \left(\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}\right), \boldsymbol {\gamma} ^ {n - k + 1} \odot \left(\boldsymbol {x} _ {k + 1} \mathbf {1} _ {c} ^ {\top}\right), \dots \boldsymbol {\gamma} ^ {0} \odot \left(\boldsymbol {x} _ {n} \mathbf {1} _ {c} ^ {\top}\right) \tag {39}
$$
However, we can factor out  $\gamma^{-p}$  from Equation 35 via a combination of the distributive property and the product of exponentials
$$
\boldsymbol {S} _ {k + p} = \gamma^ {p} \odot \left(\sum_ {j = 0} ^ {p} \boldsymbol {\gamma} ^ {- j} \odot \left(\boldsymbol {x} _ {k + j} \mathbf {1} _ {c} ^ {\top}\right)\right) + \boldsymbol {\gamma} ^ {p + 1} \odot \boldsymbol {S} _ {k - 1} \tag {40}
$$
Now, each  $\pmb{\gamma}^j$  is associated with a single  $\pmb{x}_j$ , requiring just  $n - k + 1$  terms:
$$
\boldsymbol {\gamma} ^ {0} \odot \left(\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}\right) \tag {41}
$$
$$
\boldsymbol {\gamma} ^ {0} \odot \left(\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}\right), \boldsymbol {\gamma} ^ {1} \odot \left(\boldsymbol {x} _ {k + 1} \mathbf {1} _ {c} ^ {\top}\right) \tag {42}
$$
$$
\vdots \tag {43}
$$
$$
\boldsymbol {\gamma} ^ {0} \odot \left(\boldsymbol {x} _ {k} \mathbf {1} _ {c} ^ {\top}\right), \boldsymbol {\gamma} ^ {1} \odot \left(\boldsymbol {x} _ {k + 1} \mathbf {1} _ {c} ^ {\top}\right), \dots , \boldsymbol {\gamma} ^ {n} \odot \left(\boldsymbol {x} _ {n} \mathbf {1} _ {c} ^ {\top}\right) \tag {44}
$$
We can represent each of these rows as a slice of a single  $n - k + 1$  length tensor, for a space complexity of  $O(n - k)$  or for  $k = 1$ ,  $O(n)$ .
Finally, we want to swap the exponent signs because it provides better precision when working with floating point numbers. Computing small values, then big values is more numerically stable than the other way around. We want to compute the inner sum using small numbers ( $\gamma^{+}$  results in small numbers while  $\gamma^{-}$  produces big numbers;  $n - k$  is positive and  $k - n$  is negative). We can factor  $\gamma^{k - n}$  out of the sum and rewrite Equation 40 as
$$
\boldsymbol {S} _ {k + p} = \gamma^ {k - n + p} \odot \left(\sum_ {j = 0} ^ {p} \boldsymbol {\gamma} ^ {n - k - j} \odot \left(\boldsymbol {x} _ {k + j} \mathbf {1} _ {c} ^ {\top}\right)\right) + \boldsymbol {\gamma} ^ {p + 1} \odot \boldsymbol {S} _ {k - 1} \tag {45}
$$
If we let  $t = n - k$ , this is equivalent to Equation 13:
$$
\boldsymbol {S} _ {k + p} = \boldsymbol {\gamma} ^ {p + 1} \odot \boldsymbol {S} _ {k - 1} + \boldsymbol {\gamma} ^ {p - t} \odot \sum_ {j = 0} ^ {p} \boldsymbol {\gamma} ^ {t - j} \odot \left(\boldsymbol {x} _ {k + j} \mathbf {1} _ {c} ^ {\top}\right), \quad \boldsymbol {S} _ {k + p} \in \mathbb {C} ^ {m \times c}. \tag {46}
$$
We also need to compute  $2(n - k + 1)$  gamma terms:  $\pmb{\gamma}^{k - n},\dots \pmb{\gamma}^{n - k + 1}$ , resulting in linear space complexity  $O(n)$  for a sequence of length  $n$ .
We utilize the episodic reward and normalized episodic reward metrics. We record episodic rewards as a single mean for each epoch, and report the maximum reward over a trial.
We do not modify the hyperparameters for PPO, SAC, or TD3 from the original benchmark papers. We direct the reader to [Morad et al., 2023] for PPO hyperparameters and [Ni et al., 2022] for SAC and TD3 hyperparameters.
D.1 Hardware and Efficiency Information
We trained on a server with a Xeon CPU running Torch version 1.13 with CUDA version 11.7, with consistent access to two 2080Ti GPUs. Wandb reports that we used roughly 161 GPU days of compute to produce the POMDP-Baselines results, 10 GPU days for the POPGym results, and 45 GPU days for the FFM ablation using POPGym. This results in a total of 216 GPU days. Since we ran four jobs per GPU, this corresponds to 54 days of wall-clock time for all experiments.
D.2 PPO and POPGym Baselines
Morad et al. [2023] compares models along the recurrent state size, with a recurrent state size of 256. For fairness, We let  $m = 32$  and  $c = 4$ , which results in a complex recurrent state of 128, which can be represented as a 256 dimensional real vector. We initialize  $\alpha, \omega$  following Appendix B for  $t_e = 1024, \beta = 0.01$ . We run version 0.0.2 of POPGym, and compare FFM numerically with the MMER score from the paper in Table 4.
D.3 SAC, TD3, and POMDP-Baselines