SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Then the discrete-time Fourier transform of  $\widetilde{s}_t$  is  $\mathcal{F}\widetilde{s}_t(\omega) = \sum_{n=0}^{+\infty}[\widetilde{s}_t]_n e^{-j\omega n}$ , where  $\omega$  represents the frequency variable. The discount factor  $\gamma$  in (5) is used to ensure the convergence of the Fourier transform and also serves as the contraction factor in the following Theorem 4. Since the state sequences are discrete-time signals, the corresponding DTFT is  $2\pi$ -periodic with respect to  $\omega$ . Based on this property, a common practice for operational feasibility is to compute a discrete approximation of the DTFT over one period, by sampling the DTFT at discrete points over  $[0, 2\pi]$ . In practice, we take  $L$  equally-spaced samples of the DTFT. Then the prediction target is a matrix with size  $L*D$ , where  $D$  is the dimension of the state space. We can derive that the DTFT functions at successive time steps are related to each other in a recursive form:
$$
F _ {\pi , p} \left(s _ {t}, a _ {t}\right) = \widetilde {\boldsymbol {S}} _ {t} + \Gamma E _ {\pi , p} \left[ F \left(s _ {t + 1}, a _ {t + 1}\right) \right]. \tag {6}
$$
The detailed derivation and the specific form of  $\widetilde{S}_t$  and  $\Gamma$  is provided in Appendix A.3.
Based on the recursive formula (6), we can obtain the prediction loss by computing the difference between the estimated Fourier value  $F_{\pi ,p}(s_t,a_t)$  and the better estimate  $\widetilde{S}_t + \Gamma E_{\pi ,p}[F(s_{t + 1},a_{t + 1})]$ , just like the TD error. Similar to the TD-learning of value functions, the recursive relationship can be reformulated as contraction mapping  $\mathcal{T}$ , as shown in the following theorem (see proof in Appendix A.3). Due to the properties of contraction mappings, we can iteratively apply the operator  $\mathcal{T}$  to compute the target DTFT function of long-term state sequences until convergence in tabular settings. When calculating the prediction loss, we only need to utilize the current state  $s_t$ , the current action  $a_{t}$ , and the next state  $s_{t + 1}$ . Therefore, one notable advantage of SPF is that there is no need to store multi-step future states as labels for predicting future state sequences.
Theorem 4. Let  $\mathcal{F}$  denote the set of all functions  $F: S \times \mathcal{A} \to \mathbb{C}^{L*D}$  and define the norm on  $\mathcal{F}$  as
$$
\| F\|_{\mathcal{F}}:= \sup_{\substack{s\in \mathcal{S}\\ a\in \mathcal{A}}}\max_{0\leq k <   L}\big\| \big[F(s,a)\big]_{k}\big\|_{D},
$$
where  $\left[F(s,a)\right]_k$  represents the kth row vector of  $F(s,a)$ . We show that the mapping  $\mathcal{T}:\mathcal{F}\to \mathcal{F}$  defined as
$$
\mathcal {T} F \left(s _ {t}, a _ {t}\right) = \widetilde {\boldsymbol {S}} _ {t} + \Gamma E _ {\pi , p} [ F \left(s _ {t + 1}, a _ {t + 1}\right) ] \tag {7}
$$
is a contraction mapping, where  $\widetilde{S}_t$  and  $\Gamma$  are defined in Appendix A.3.
As the Fourier transform of the real state signals has the property of conjugate symmetry, we only need to predict the DTFT on a half-period interval  $[0,\pi ]$ . Therefore, we reduce the row size of the
Figure 2: The network architecture of SPF. The online encoder  $\phi$  outputs the representations used in the RL task and the predictor  $\mathcal{F}$  predicts complex-valued Fourier transform of the state sequences starting from the state-action pair  $(s_t, a_t)$ . During training,  $(s_t, a_t, s_{t+1})$  are previously experienced states and actions sampled from a replay buffer. The dashed line shows how gradients flow back to model weights. We prevent the gradient of RL losses from updating the online encoder and prevent the gradient of prediction loss from updating the target encoder.
prediction target to half for reducing redundant information and saving storage space. In practice, we train a parameterized prediction model  $\mathcal{F}$  to predict the DTFT of state sequences. Note that the value of the prediction target is on the complex plane, so the prediction network employs two separate output modules  $\mathcal{F}_{\mathrm{Re}}$  and  $\mathcal{F}_{\mathrm{Im}}$  as real and imaginary parts respectively. Then we define the auxiliary prediction loss function as:
$$
\begin{array}{l} L _ {\text {p r e d}} (\phi , \mathcal {F}) = d \left(\widetilde {\boldsymbol {S}} _ {t} + \left[ \Gamma_ {\text {R e}} \mathcal {F} _ {\text {R e}} \left(\overline {{s _ {t + 1}}}, \pi \left(\overline {{s _ {t + 1}}}\right)\right) - \Gamma_ {\text {I m}} \mathcal {F} _ {\text {I m}} \left(\overline {{s _ {t + 1}}}, \pi \left(\overline {{s _ {t + 1}}}\right)\right) \right], \mathcal {F} _ {\text {R e}} \left(\overline {{s _ {t}}}, a _ {t}\right)\right) \\ + d \Big (\left[ \Gamma_ {\operatorname {I m}} \mathcal {F} _ {\operatorname {R e}} \left(\overline {{s _ {t + 1}}}, \pi \left(\overline {{s _ {t + 1}}}\right)\right) + \Gamma_ {\operatorname {R e}} \mathcal {F} _ {\operatorname {I m}} \left(\overline {{s _ {t + 1}}}, \pi \left(\overline {{s _ {t + 1}}}\right)\right) \right], \mathcal {F} _ {\operatorname {I m}} \left(\overline {{s _ {t}}}, a _ {t}\right) \Big), \\ \end{array}
$$
where  $\overline{s_t} = \phi(s_t)$  means the representation of the state,  $\Gamma_{\mathrm{Re}}$  and  $\Gamma_{\mathrm{Im}}$  denote the real and imaginary parts of the constant  $\Gamma$ , and  $d$  denotes an arbitrary similarity measure. We choose  $d$  as cosine similarity in practice. The algorithm pseudo-code is shown in Appendix B.
5.3 Network Architecture of SPF
Here we provide more details about the network architecture of our method. In addition to the encoder and the predictor, we use projection heads  $\psi$  [25] that project the predicted values onto a low-dimensional space to prevent overfitting when computing the prediction loss directly from the high-dimensional predicted values. In practice, we use a loss function called freqloss, which preserves the low and high-frequency components of the predicted DTFT without the dimensionality reduction process (See Appendix C for more details). Furthermore, when computing the target predicted value, we follow prior work [26, 11] to use the target encoder, predictor, and projection for more stable performance. We periodically overwrite the target network parameters with an exponential moving average of the online network parameters.
In the training process, we train the encoder  $\phi$ , the predictor  $\mathcal{F}$ , and the projection  $\psi$  to minimize the auxiliary prediction loss  $\mathcal{L}_{\mathrm{pred}}(\phi, \mathcal{F}, \psi)$ , and alternately update the actor-critic models of RL tasks using the trained encoder  $\phi$ . We illustrate the overall architecture of SPF and the gradient flows during training in Figure 2.
We quantitatively evaluate our method on a standard continuous control benchmark—the set of MuJoCo [27] environments implemented in OpenAI Gym.
6.1 Comparative Evaluation on MuJoCo
To evaluate the effect of learned representations, we measure the performance of two traditional RL methods SAC [28] and PPO [29] with raw states, OFENet representations [19], and SPF represent-
Table 1: Mean and standard error results on six MuJoCo tasks at 500K step and 1M step. All means and standard errors are calculated over 10 seeds. The highest mean scores are marked in blue. SPF outperforms other SOTA methods in 5 out of 6 settings with an average  $19.5\%$  boost.
Environment	SAC-SPF	SAC-OFE	SAC-raw	PPO-SPF	PPO-OFE	PPO-raw
500K step	HalfCheetah	12739 ± 139 (+7%)	11932 ± 75	9004 ± 150	2448 ± 357 (+1%)	2419 ± 296	1475 ± 238
Hopper	3443 ± 24 (+15%)	2983 ± 123	2343 ± 183	2923 ± 98 (+36%)	1529 ± 157	2156 ± 245
Walker2d	4868 ± 53 (+29%)	3762 ± 162	2023 ± 269	1203 ± 211 (+45%)	339 ± 32	827 ± 156
Ant	6531 ± 175 (+17%)	5587 ± 253	2691 ± 76	867 ± 34	696 ± 47	929 ± 15
Swimmer	45 ± 0	45 ± 0	42 ± 0	90 ± 11 (+15%)	78 ± 12	66 ± 6
Humanoid	4613 ± 118 (+13%)	4071 ± 172	2953 ± 136	395 ± 10	405 ± 18	383 ± 9
1M step	HalfCheetah	15822 ± 109 (+10%)	14425 ± 112	10745 ± 159	3154 ± 1358 (+3%)	3066 ± 326	2259 ± 344
Hopper	3517 ± 20 (+10%)	3197 ± 147	3056 ± 91	3152 ± 137 (+16%)	2370 ± 258	2721 ± 222
Walker2d	5158 ± 69 (+7%)	4833 ± 59	3367 ± 110	2229 ± 299	1080 ± 236	2302 ± 158
Ant	7241 ± 85 (+7%)	6738 ± 150	4220 ± 164	1053 ± 79 (+15%)	913 ± 18	867 ± 17
Swimmer	45 ± 0	45 ± 0	43 ± 0	99 ± 13 (+19%)	73 ± 10	83 ± 9
Humanoid	5633 ± 112	6241 ± 98	4618 ± 118	441 ± 17	448 ± 17	439 ± 18
Figure 3: Results on six MuJoCo tasks. The solid curves denote the means and the shaded regions denote the minimum and maximum returns over 10 seeds. Each checkpoint is evaluated by 10 episodes in evaluated environments. Curves are smoothed for visual clarity.
tations. We train our representations via the Fourier transform of infinite-horizon state sequences, which means the representations learn to capture the information of infinite-step future states. Hence, we compare our representation with OFENet, which uses the auxiliary task of predicting single-step future states for representation learning. For a fair comparison, we use an encoder with the same network structure as that of OFENet. See Appendix C and D for more details about network architectures and hyperparameters setting.
As described above, our comparable methods include: 1) SAC-SPF: SAC with SPF representations; 2) SAC-OFE: SAC with OFENet representations; 3) SAC-raw: SAC with raw states; 4) PPO-SPF: PPO with SPF representations; 5) PPO-OFE: PPO with OFENet representations; 6) PPO-raw: PPO with raw states. Figure 3 shows the learning curves of the above methods. As our results suggest, SPF shows superior performance compared to the original algorithms, SAC-raw and PPO-raw, across all six MuJoCo tasks and also outperforms OFENet in terms of both sample efficiency and asymptotic performance on five out of six MuJoCo tasks. According to the results in Table 1, SPF achieves an average gain of  $+19.5\%$  over other methods at 500K step and 1M step. These results indicate that our approach learns more quickly than other methods when the number of interactions is limited. One possible explanation for the comparatively weaker performance of our method on the Humanoid-v2
(b)
Figure 4: Results of additional trials. (a) Ablation study: training curves of five variant methods of SPF on SAC. Each method is evaluated over 5 seeds on HalfCheetah-v2. (b) Comparison of prediction targets: training curves of SAC with different auxiliary prediction tasks. Each method is evaluated over 5 seeds on HalfCheetah-v2. (c) Visualization of recovered states from the predicted DTFT. The blue line represents the true state sequence, while the red line represents the recovered state sequence. The lighter red line corresponds to predictions made by historical states from a more distant time step. The four subfigures represent four dimensions of the state space on Walker2d.
task is that its state contains the external forces to bodies, which show limited regularities due to their discontinuous and sparse.
In this part, we will verify that just predicting the FT of state sequences may fall short of the expected performance and that using SPF is necessary to get better performance. To this end, we conducted an ablation study to identify the specific components that contribute to the performance improvements achieved by SPF. Figure 4(a) shows the ablation study over SAC with HalfCheetah-v2 environment.
notarg removes all target networks of the encoder, predictor, and projection layer from SPF. Based on the results, the variant of SPF exhibits significantly reduced performance when target estimations in the auxiliary loss are generated by the online encoder without a stopgradient. Therefore, using a separate target encoder is vital, which can significantly improve the stability and convergence properties of our algorithm.
nofreqloss computes the cosine similarity of the projection layer's outputs directly without any special treatment to the low-frequency and high-frequency components of our predicted DTFT. The reduced convergence rate of nofreqloss suggests that preserving the complete information of low and high-frequency components can encourage the representations to capture more structural information in the frequency domain.
noproj removes the projection layer from SPF and computes the cosine similarity of the predicted values as the objective. The performance did not significantly deteriorate after removing the projection layer, which indicates that the prediction accuracy of state sequences in the frequency domain may not have a strong impact on the quality of representation learning. Therefore, it can be inferred that SPF places a greater emphasis on capturing the underlying structural information of the state sequences, and is capable of reconstructing the state sequences with a low risk of overfitting.
mlp changes the lay block of the encoder from MLP-DenseNet to MLP. The much lower scores of mlp indicate that both the raw state and the output of hidden layers contain important information that contributes to the quality of the learned representations. This result underscores the importance of leveraging sufficient information for representation learning.
mlp-cat uses a modified block of MLP as the layer of the encoder, which concat the output of MLP with the raw state. The performance of mlp-cat does increase compared to mlp, but is still not as good as SPF in terms of both sample efficiency and performance.
6.3 Comparison of Different Prediction Targets
This section aims to test the effect of our prediction target—infinite-step state sequences in the frequency domain—on the efficiency of representation learning. We test five types of prediction
targets: 1) Sta1: single-step future state; 2) StaN: N-step state sequences, where we choose  $N = 2,3,5;3$  SPF: infinite-step state sequences in frequency domain; 4) Rew1: single-step future reward; 5) RewFT: infinite-step reward sequences in frequency domain.
As shown in Figure 4(b), SPF outperforms all other competitors in terms of sample efficiency, which indicates that infinite-step state sequences in the frequency domain contain more underlying valuable information that can facilitate efficient representation learning. Since Sta1 and SPF outperform Rew1 and RewFT respectively, it can be referred that learning via states is more effective for representation learning than learning via rewards. Notably, the lower performance of StaN compared to Sta1 could be attributed to the model's tendency to prioritize prediction accuracy over capturing the underlying structured information in the sequential data, which may impede its overall learning efficiency.
6.4 Visualization of Recovered State Sequences
This section aims to demonstrate that the representations learned by SPF effectively capture the structural information contained in infinite-step state sequences. To this end, we compare the true state sequences with the states recovered from the predicted DTFT via the inverse DTFT (See Appendix E for more implementation details). Figure 4(c) shows that the learned representations can recover the true state sequences even using the historical states that are far from the current time step. In Appendix E, we also provide a visualization of the predicted DTFT, which is less accurate than the results in Figure 4(c). Those results highlight the ability of SPF to effectively extract the underlying structural information in infinite-step state sequences without relying on high prediction accuracy.
We further provide a comparison table that measures the distance between the real DTFT and the predicted DTFT using cosine similarity. The results provided in Appendix F indicate that the prediction module  $\mathcal{F}$  exhibits moderate predictive accuracy in approximating the real Fourier transform, with an average cosine similarity value of  $-0.6$ .
In this paper, we theoretically analyzed the existence of structural information in state sequences, which is closely related to policy performance and signal regularity, and then introduced State Sequences Prediction via Fourier Transform (SPF), a representation learning method that predicts the FT of state sequences to extract the underlying structural information in state sequences for learning expressive representations efficiently. SPF outperforms several state-of-the-art algorithms in terms of both sample efficiency and performance. Our additional experiments and visualization show that SPF encourages representations to place a greater emphasis on capturing the underlying pattern of time-series data, rather than pursuing high accuracy of prediction tasks.
Limitations One of the main limitations of our paper is that we have only evaluated our method on tasks where the state sequences exhibit strong asymptotic periodicity. Considering that the Fourier transform converts non-periodic signals into a continuous frequency domain, it is more difficult to extract meaningful frequency features of non-periodic signals compared to periodic signals that have a discrete frequency domain. Moreover, the frequency features extracted by our approach inherently depend on the policy and task, which limits their reusability across multiple tasks. Further research is needed to analyze the applicability of our approach to non-periodic tasks and the potential for generalization across multiple tasks.
The authors would like to thank all the anonymous reviewers for their insightful comments. This work was supported by National Key R&D Program of China under contract 2022ZD0119801, National Nature Science Foundations of China grants U19B2026, U19B2044, 61836011, 62021001, and 61836006.
[1] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533, 2015.
[20] Sridhar Mahadevan. Proto-value functions: Developmental reinforcement learning. In Proceedings of the 22nd international conference on Machine learning, pages 553-560, 2005.
[21] Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. Wiley Series in Probability and Statistics. Wiley, 1994.
[22] Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pages 22-31, 2017.
[23] Ofir Nachum and Mengjiao Yang. Provable representation learning for imitation with contrastive fourier features. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, pages 30100-30112, 2021.
[24] Sham M. Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Machine Learning, Proceedings of the Nineteenth International Conference, pages 267-274, 2002.
[25] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, ICML, volume 119, pages 1597-1607, 2020.
[26] Jean-Bastien Grill, Florian Strub, Florent Alché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent - A new approach to self-supervised learning. volume 33, pages 21271-21284, 2020.
[27] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems, IROS, pages 5026-5033. IEEE, 2012.
[28] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, ICML, volume 80, pages 1856-1865, 2018.
[29] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, 2017.
[30] Geoffrey Grimmett and David Stirzaker. Probability and random processes. Oxford university press, 2020.
[31] Eugene Seneta. Non-negative matrices and Markov chains. Springer Science & Business Media, 2006.
[32] Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Proceedings of the 35th International Conference on Machine Learning, ICML, volume 80 of Proceedings of Machine Learning Research, pages 1582-1591, 2018.
A.1 Proof of Performance Difference Distinction via State Sequences
Following the previous work [22], our analysis will make use of the discounted future state distribution,  $d^{\pi}$ , which is defined as
$$
d ^ {\pi} (s) = (1 - \gamma) \sum_ {t = 0} ^ {\infty} \gamma^ {t} P (s _ {t} = s | \pi , \mathcal {M})
$$
It allows us to express the expected discounted total reward compactly as
$$
\begin{array}{l} J (\pi) = \sum_ {t = 0} ^ {\infty} \gamma^ {t} E _ {s _ {t}, a _ {t}, s _ {t + 1}} [ R (s _ {t}, a _ {t}, s _ {t + 1}) | \pi , \mathcal {M} ] \\ = \sum_ {t = 0} ^ {\infty} \gamma^ {t} \int_ {\mathcal {S}} R _ {\pi} (s) P (s _ {t} = s | \pi , \mathcal {M}) \mathrm {d} s \\ = \int_ {\mathcal {S}} R _ {\pi} (s) \sum_ {t = 0} ^ {\infty} \gamma^ {t} P \left(s _ {t} = s \mid \pi , \mathcal {M}\right) \mathrm {d} s (1) \\ = \frac {1}{1 - \gamma} \int_ {S} R _ {\pi} (s) d _ {\pi} (s) \mathrm {d} s \\ = \frac {1}{1 - \gamma} E _ {\substack {s \sim d ^ {\pi} \\ a \sim \pi \\ s ^ {\prime} \sim P}} \left[ R (s, a, s ^ {\prime}) \right], (2) \\ \end{array}
$$
where we define  $R_{\pi}(s) \coloneqq E_{a \sim \pi, s' \sim P}[R(s, a, s')]$ . It should be clear from  $a \sim \pi(\cdot | s)$  and  $s' \sim P(\cdot | s, a)$  that  $a$  and  $s'$  depend on  $s$ . Thus, the reward function  $R_{\pi}$  is only related to  $s$  when the policy  $\pi$  is fixed.
Firstly, we prove that the distance between two state sequence distributions obtained from two distinct policies serves as an upper bound on the performance difference between those policies, provided that certain assumptions regarding the reward function hold.
Proof. According to the equation (1), the difference in performance between two policies  $\pi_1, \pi_2$  can be bounded as follows.
$$
\begin{array}{l} | J (\pi_ {1}) - J (\pi_ {2}) | \leq R _ {\max } \cdot \sum_ {t = 0} ^ {\infty} \gamma^ {t} \int_ {\mathcal {S}} \left| P (\mathbf {s} _ {\mathbf {t}} = s | \pi_ {1}, \mathcal {M}) - P (\mathbf {s} _ {\mathbf {t}} = s | \pi_ {2}, \mathcal {M}) \right| d s \\ \leq R _ {\max } \cdot \sum_ {t = 0} ^ {T} \gamma^ {t} \int_ {\mathcal {S}} \left| \int_ {\mathcal {S} ^ {T}} P \left(s _ {0}, \dots , s _ {t - 1}, s, s _ {t + 1}, \dots , s _ {T} \mid \pi_ {1}, \mathcal {M}\right) \right. \\ - \left. P (s _ {0}, \ldots , s _ {t - 1}, s, s _ {t + 1} \ldots , s _ {T} | \pi_ {2}, \mathcal {M}) \mathrm {d} s _ {0} \dots \mathrm {d} s _ {t - 1} \mathrm {d} s _ {t + 1} \dots \mathrm {d} s _ {T} \right| \mathrm {d} s \\ + R _ {\max } \cdot 2 \sum_ {t = T + 1} ^ {\infty} \gamma^ {t}, \quad \forall T \geq 1 \\ \end{array}
$$
$$
\begin{array}{l} \leq R _ {\max } \sum_ {t = 0} ^ {T} \gamma^ {t} \int_ {\mathcal {S} ^ {T + 1}} \left| P \left(s _ {0}, \dots , s _ {T} \mid \pi_ {1}, \mathcal {M}\right) - P \left(s _ {0} \dots , s _ {T} \mid \pi_ {2}, \mathcal {M}\right) \right| d s _ {0} \dots d s _ {T} \\ + R _ {\max } \cdot 2 \sum_ {t = T + 1} ^ {\infty} \gamma^ {t}, \quad \forall T \geq 1 \\ = \frac {R _ {\operatorname* {m a x}}}{1 - \gamma} \cdot \int_ {\mathcal {S} ^ {T + 1}} \left| P \left(s _ {0}, \dots , s _ {T} \mid \pi_ {1}, \mathcal {M}\right) - P \left(s _ {0} \dots , s _ {T} \mid \pi_ {2}, \mathcal {M}\right) \right| d s _ {0} \dots d s _ {T} \\ + R _ {\max } \cdot 2 \sum_ {t = T + 1} ^ {\infty} \gamma^ {t}, \quad \forall T \geq 1. \\ \end{array}
$$
Let  $T\to \infty$  , then we obtain the bound proposed by (3).