SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mingxuan Ye
Yufei Kuang
Jie Wang $^{1,2*}$
Rui Yang
Wengang Zhou<sup>1,2</sup>
Houqiang  $\mathbf{L}\mathbf{i}^{1,2}$
Feng Wu
$^{1}$ CAS Key Laboratory of Technology in GIPAS, University of Science and Technology of China
$^{2}$ Institute of Artificial Intelligence, Hefei Comprehensive National Science Center
{mingxuanye, yfkuang, yr0013}@mail.ustc.edu.cn
{jiewangx, zhwg, lihq, fengwu}@ustc.edu.cn
While deep reinforcement learning (RL) has been demonstrated effective in solving complex control tasks, sample efficiency remains a key challenge due to the large amounts of data required for remarkable performance. Existing research explores the application of representation learning for data-efficient RL, e.g., learning predictive representations by predicting long-term future states. However, many existing methods do not fully exploit the structural information inherent in sequential state signals, which can potentially improve the quality of long-term decision-making but is difficult to discern in the time domain. To tackle this problem, we propose State Sequences Prediction via Fourier Transform (SPF), a novel method that exploits the frequency domain of state sequences to extract the underlying patterns in time series data for learning expressive representations efficiently. Specifically, we theoretically analyze the existence of structural information in state sequences, which is closely related to policy performance and signal regularity, and then propose to predict the Fourier transform of infinite-step future state sequences to extract such information. One of the appealing features of SPF is that it is simple to implement while not requiring storage of infinite-step future states as prediction targets. Experiments demonstrate that the proposed method outperforms several state-of-the-art algorithms in terms of both sample efficiency and performance.[2]
Deep reinforcement learning (RL) has achieved remarkable success in complex sequential decision-making tasks, such as computer games [1], robotic control [2], and combinatorial optimization [3]. However, these methods typically require large amounts of training data to learn good control, which limits the applicability of RL algorithms to real-world problems. The crucial challenge is to improve the sample efficiency of RL methods. To address this challenge, previous research has focused on representation learning to extract adequate and valuable information from raw sensory data and train RL agents in the learned representation space, which has been shown to be significantly more data-efficient [4-8]. Many of these algorithms rely on auxiliary self-supervision tasks, such as predicting future reward signals [4] and reconstructing future observations [5], to incorporate prior knowledge about the environment into the representations.
2The code of SPF is available on GitHub at https://github.com/MIRALab-USTC/RL-SPF/.
Due to the sequential nature of RL tasks, multi-step future signals inherently contain more features that are valuable for long-term decision-making than immediate future signals. Recent work has demonstrated that leveraging future reward sequences as supervisory signals is effective in improving the generalization performance of visual RL algorithms [9, 10]. However, we argue that the state sequence provides a more informative supervisory signal compared to the sparse reward signal. As shown in the top portion of Figure 1, the sequence of future states essentially determines future actions and further influences the sequence of future rewards. Therefore, state sequences maximally preserve the influence of the transition intrinsic to the environment and the effect of actions generated from the current policy.
Despite these benefits of state sequences, it is challenging to extract features from these sequential data through long-term prediction tasks. A substantial obstacle is the difficulty of learning accurate long-term
prediction models for feature extraction. Previous methods propose making multi-step predictions using a one-step dynamic model by repeatedly feeding the prediction back into the learned model [11- 13]. However, these approaches require a high degree of accuracy in the one-step model to avoid accumulating errors in multi-step predictions [13]. Another obstacle is the storage of multi-step prediction targets. For instance, some method learns a specific dynamics model to directly predict multi-step future states [14], which requires significant additional memory to store multi-step future states as prediction targets.
To tackle these problems, we propose utilizing the structural information inherent in the sequential state signals to extract useful features, thus circumventing the difficulty of learning an accurate prediction model. In Section 4, we theoretically demonstrate two types of sequential dependency structures present in state sequences. The first structure involves the dependency between reward sequences and state sequences, where the state sequences implicitly reflect the performance of the current policy and exhibit significant differences under good and bad policies. The second structure pertains to the temporal dependencies among the state signals, namely the regularity patterns exhibited by the state sequences. By exploiting the structural information, representations can focus on the underlying critical features of long-term signals, thereby reducing the need for high prediction accuracy and improving training stability [13].
Building upon our theoretical analyses, we propose State Sequences Prediction via Fourier Transform (SPF), a novel method that exploits the frequency domain of state sequences to efficiently extract the underlying structural information of long-term signals. Utilizing the frequency domain offers several advantages. Firstly, it is widely accepted that the frequency domain shows the regularity properties of the time-series data [15-17]. Secondly, we demonstrate in Section 5.1 that the Fourier transform of state sequences retains the ability to indicate policy performance under certain assumptions. Moreover, Figure 1 provides an intuitive understanding that the frequency domain enables more effective discrimination of two similar temporal signals that are difficult to differentiate in the time domain, thereby improving the efficiency of policy performance distinction and policy learning.
Specifically, our method performs an auxiliary self-supervision task that predicts the Fourier transform (FT) of infinite-step state sequences to improve the efficiency of representation learning. To facilitate the practical implementation of our method, we reformulate the Fourier transform of state sequences as a recursive form, allowing the auxiliary loss to take the form of a TD error [18], which depends only on the single-step future state. Therefore, SPF is simple to implement and eliminates the requirements for storing infinite-step state sequences when computing the labels of FT. Experiments demonstrate that our method outperforms several state-of-the-art algorithms in terms of both sample efficiency and performance on six MuJoCo tasks. Additionally, we visualize the fine distinctions between the multi-step future states recovered from our predicted FT and the true states, which indicates that our representation effectively captures the inherent structures of future state sequences.
Figure 1: State Sequence Generation Process. The top row displays a sample of state sequences generated from MDP  $\mathcal{M}$ . The bottom two rows visualize two state sequences in the time and frequency domains respectively. Each column corresponds to a state sequence generated from a different policy.
Learning Predictive Representations in RL. Many existing methods leverage auxiliary tasks that predict single-step future reward or state signals to improve the efficiency of representation learning [2, 5, 19]. However, multi-step future signals inherently contain more valuable features for long-term decision-making than single-step signals. Recent work has demonstrated the effectiveness of using future reward sequences as supervisory signals to improve the generalization performance of visual RL algorithms [9, 10]. Several studies propose making multi-step predictions of state sequences using a one-step dynamic model by repeatedly feeding the prediction back into the learned model, which is applicable to both model-free [11] and model-based [12] RL. However, these approaches require a high degree of accuracy in the one-step model to prevent accumulating errors [13]. To tackle this problem, the existing method learned a prediction model to directly predict multi-step future states [14], which results in significant additional storage for the prediction labels. In our work, we propose to predict the FT of state sequences, which reduces the demand for high prediction accuracy and eliminates the need to store multi-step future states as prediction targets.
Incorporating the Fourier Features. There existing many traditional RL methods to express the Fourier features. Early works have investigated representations based on a fixed basis such as Fourier basis to decompose the function into a sum of simpler periodic functions [20, 18]. Another research explored enriching the representational capacity using random Fourier features of the observations. Moreover, in the field of self-supervised pre-training, neuro2vec [16] conducts representation learning by predicting the Fourier transform of the masked part of the input signal, which is similar to our work but needs to store the entire signal as the label.
3.1 MDP Notation
In RL tasks, the interaction between the environment and the agent is modeled as a Markov decision process (MDP). We consider the standard MDP framework [21], in which the environment is given by the tuple  $\mathcal{M} \coloneqq \langle S, \mathcal{A}, R, P, \mu, \gamma \rangle$ , where  $\mathcal{S}$  is the set of states,  $\mathcal{A}$  is the set of actions,  $R: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \to [-R_{\max}, R_{\max}]$  is the reward function,  $P: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \to [0, 1]$  is the transition probability function,  $\mu: \mathcal{S} \to [0, 1]$  is the initial state distribution, and  $\gamma \in [0, 1)$  is the discount factor. A policy  $\pi$  defines a probability distribution over actions conditioned on the state, i.e.  $\pi(a|s)$ . The environment starts at an initial state  $s_0 \sim \mu$ . At time  $t \geq 0$ , the agent follows a policy  $\pi$  and selects an action  $a_t \sim \pi(\cdot | s_t)$ . The environment then stochastically transitions to a state  $s_{t+1} \sim P(\cdot | s_t, a_t)$  and produces a reward  $r_t = R(s_t, a_t, s_{t+1})$ . The goal of RL is to select an optimal policy  $\pi^*$  that maximizes the cumulative sum of future rewards. Following previous work [22, 23], we define the performance of a policy  $\pi$  as its expected sum of future discounted rewards:
$$
J (\pi , \mathcal {M}) := E _ {\tau \sim (\pi , \mathcal {M})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} R \left(s _ {t}, a _ {t}, s _ {t + 1}\right) \right], \tag {1}
$$
where  $\tau := (s_0, a_0, s_1, a_1, \dots)$  denotes a trajectory generated from the interaction process and  $\tau \sim (\pi, \mathcal{M})$  indicates that the distribution of  $\tau$  depends on  $\pi$  and the environment model  $\mathcal{M}$ . For simplicity, we write  $J(\pi)$  and  $\tau \sim \pi$  as shorthand since our environment is stationary. We also interest about the discounted future state distribution  $d^{\pi}$ , which is defined by  $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
$$
J (\pi) = \frac {1}{1 - \gamma} \underset { \begin{array}{c} a \sim \pi \\ s ^ {\prime} \sim P \end{array} } {E _ {s \sim d ^ {\pi}}} \left[ R (s, a, s ^ {\prime}) \right]. \tag {2}
$$
The proof of (2) can be found in [24] or Section A.1 in the appendix.
Given that we aim to train an encoder that effectively captures the useful aspects of the environment, we also consider a latent MDP  $\overline{\mathcal{M}} = \langle \overline{S},\mathcal{A},\overline{R},\overline{P},\overline{\mu},\gamma ,\rangle$  , where  $\overline{S}\subset \mathbb{R}^D$  for finite  $D$  and the action space  $\mathcal{A}$  is shared by  $\mathcal{M}$  and  $\overline{\mathcal{M}}$  . We aim to learn an embedding function  $\phi :S\to \overline{S}$  , which connects the state spaces of these two MDPs. We similarly denote  $\overline{\pi}^{*}$  as the optimal policy in  $\overline{\mathcal{M}}$  . For ease of notation, we use  $\overline{\pi} (\cdot |s)\coloneqq \overline{\pi} (\cdot |\phi (s))$  to represent first using  $\phi$  to map  $s\in S$  to the latent state space  $\overline{S}$  and subsequently using  $\overline{\pi}$  to generate the probability distribution over actions.
3.2 Discrete-Time Fourier Transform
The discrete-time Fourier transform (DTFT) is a powerful way to decompose a time-domain signal into different frequency components. It converts a real or complex sequence  $\{x_{n}\}_{n = -\infty}^{+\infty}$  into a complex-valued function  $F(\omega) = \sum_{n = -\infty}^{\infty}x_ne^{-j\omega n}$ , where  $\omega$  is a frequency variable. Due to the discrete-time nature of the original signal, the DTFT is 2π-periodic with respect to its frequency variable, i.e.,  $F(\omega +2\pi) = F(\omega)$ . Therefore, all of our interest lies in the range  $\omega \in [0,2\pi]$  that contains all the necessary information of the infinite-horizon time series.
It is important to note that the signals considered in this paper, namely state sequences, are real-valued, which ensures the conjugate symmetry property of DTFT, i.e.,  $F(2\pi - \omega) = F^{*}(\omega)$ . Therefore, in practice, it suffices to predict the DTFT only on the range of  $[0, \pi]$ , which can reduce the number of parameters and further save storage space.
4 Structural Information in State Sequences
In this section, we theoretically demonstrate the existence of the structural information inherently in the state sequences. We argue that there are two types of sequential dependency structures present in state sequences, which are useful for indicating policy performance and capturing regularity features of the states, respectively.
4.1 Policy Performance Distinction via State Sequences
In the RL setting, it is widely recognized that pursuing the highest reward at each time step in a greedy manner does not guarantee the maximum long-term benefit. For that reason, RL algorithms optimize the objective of cumulative reward over an episode, rather than the immediate reward, to encourage the agent to make farsighted decisions. This is why previous work has leveraged information about future reward sequences to capture long-term features for stronger representation learning [9].
Compared to the sparse reward signals, we claim that sequential state signals contain richer information. In MDP, the stochasticity of a trajectory derives from random actions selected by the agent and the environment's subsequent transitions to the next state and reward. These two sources of stochasticity are modeled as the policy  $\pi(a|s)$  and the transition  $p(s', r|s, a)$ , respectively. Both of them are conditioned on the current state. Over long interaction periods, the dependencies of action and reward sequences on state sequences become more evident. That is, the sequence of future states largely determines the sequence of actions that the agent selects and further determines the corresponding sequence of rewards, which implies the trend and performance of the current policy, respectively. Thus, state sequences not only explicitly contain information about the environment's dynamics model, but also implicitly reveal information about policy performance.
We provide further theoretical justification for the above statement, which shows that the distribution distance between two state sequences obtained from different policies provides an upper bound on the performance difference between those policies, under certain assumptions about the reward function.
Theorem 1. Suppose that the reward function  $R(s, a, s') = R(s)$  is related to the state  $s$ , then the performance difference between two arbitrary policies  $\pi_1$  and  $\pi_2$  is bounded by the L1 norm of the difference between their state sequence distributions:
$$
\left| J \left(\pi_ {1}\right) - J \left(\pi_ {2}\right) \right| \leq \frac {R _ {\max }}{1 - \gamma} \cdot \| P \left(s _ {0}, s _ {1}, s _ {2}, \dots \mid \pi_ {1}, \mathcal {M}\right) - P \left(s _ {0}, s _ {1}, s _ {2}, \dots \mid \pi_ {2}, \mathcal {M}\right) \| _ {1}, \tag {3}
$$
where  $P(s_0,s_1,s_2,\ldots |\pi ,\mathcal{M})$  means the joint distribution of the infinite-horizon state sequence  $\mathbf{S} = \{\mathbf{s}_0,\mathbf{s}_1,\mathbf{s}_2,\dots \}$  conditioned on the policy  $\pi$  and the environment model  $\mathcal{M}$ .
The proof of this theorem is provided in Appendix A.1. The theorem demonstrates that the greater the difference in policy performance, the greater the difference in their corresponding state sequence distributions. When we adjust the ratio  $\frac{R_{max}}{1 - \gamma}$  to take a relatively small value by scaling the reward, the theorem indicates that good and bad policies generate significantly different state sequence distributions. Furthermore, it confirms that learning via state sequences can significantly influence the search for policies with good performance.
4.2 Asymptotic Periodicity of States in MDP
Many tasks in real scenarios exhibit periodic behavior as the underlying dynamics of the environment are inherently periodic, such as industrial robots, car driving in specific scenarios, and area sweeping tasks. Take the assembly robot as an example, the robot is trained to assemble parts together to create a final product. When the robot reaches a stable policy, it executes a periodic sequence of movements that allow it to efficiently assemble the parts together. In the case of MuJoCo tasks, the agent also exhibits periodic locomotion when reaching a stable policy. We provide a corresponding video in the supplementary material to show the periodic locomotion of several MuJoCo tasks.
Inspired by these cases, we provide some theoretical analysis to demonstrate that, under some assumptions about the transition probability matrices, the state sequences in finite state space may exhibit asymptotically periodic behaviors when the agent reaches a stable policy.
Theorem 2. Suppose that the state space  $S$  is finite with a transition probability matrix  $P \in \mathbb{R}^{|S| \times |S|}$  and  $S$  has  $\alpha$  recurrent classes. Let  $R_1, R_2, \ldots, R_\alpha$  be the probability submatrices corresponding to the recurrent classes and let  $d_1, d_2, \ldots, d_\alpha$  be the number of the eigenvalues of modulus 1 that the submatrices  $R_1, R_2, \ldots, R_\alpha$  has. Then for any initial distribution  $\mu_0$ ,  $P^n \mu_0$  is asymptotically periodic with period  $d = \operatorname{lcm}(d_1, d_2, \ldots, d_\alpha)$ .
The proof of the theorem is provided in Appendix A.2. The above theorem demonstrates that there exist regular and highly-structured features in the state sequences, which can be used to learn an expressive representation. Note that in an infinite state space, if the Markov chain contains a recurrent class, then after a sufficient number of steps, the state will inevitably enter one of the recurrent classes. In this scenario, the asymptotic periodicity of the state sequences can be also analyzed using the aforementioned theorem. Furthermore, even if the state sequences do not exhibit a strictly periodic pattern, regularities still exist within the sequential data that can be extracted as representations to facilitate policy learning.
In the previous section, we demonstrate that state sequences contain rich structural information which implicitly indicates the policy performance and regular behavior of states. However, such information is not explicitly shown in state sequences in the time domain. In this section, we describe how to effectively leverage the inherent structural information in time-series data.
5.1 Learning via Frequency Domain of State Sequences
In this part, we will discuss the advantages of leveraging the frequency pattern of state sequences for capturing the inherent structural information above explicitly and efficiently.
Based on the following theorem, we find that the FT of the state sequences preserves the property in the time domain that the distribution difference between state sequences controls the performance difference between the corresponding two policies, but is subject to some stronger assumptions.
Theorem 3. Suppose that  $\mathcal{S} \subset \mathbb{R}^D$  the reward function  $R(s, a, s') = R(s)$  is an nth-degree polynomial function with respect to  $s \in S$ , then for any two policies  $\pi_1$  and  $\pi_2$ , their performance difference can be bounded as follows:
$$
| J (\pi_ {1}) - J (\pi_ {2}) | \leq \frac {\sqrt {D}}{1 - \gamma} \cdot \sum_ {k = 1} ^ {n} \frac {\left\| R ^ {(k)} (0) \right\| _ {D}}{k !} \cdot \max  _ {1 \leq i \leq D} \sup  _ {\omega_ {i} \in [ 0, 2 \pi ]} \left| F _ {\pi_ {1}} ^ {(k)} (\omega_ {i}) - F _ {\pi_ {2}} ^ {(k)} (\omega_ {i}) \right|, \tag {4}
$$
where  $F_{\pi}^{(k)}(\omega)$  denotes the DTFT of the time series  $\mathbf{S}^{(k)} = \{\mathbf{s_0}^k,\mathbf{s_1}^k,\mathbf{s_2}^k,\dots \}$  for any integer  $k\in [1,n]$  and  $\mathbf{S}^{(k)}$  means the kth power of the state sequence produced by the policy  $\pi$ . The dimensionality of  $\omega$  is the same as  $s$ .
We provide the proof in Appendix A.1. Similar to the analysis of Theorem 1, the above theorem shows that state sequences in the frequency domain can indicate the policy performance and can be leveraged to enhance the search for optimal policies. Furthermore, the Fourier transform can decompose the state sequence signal into multiple physically meaningful components. This operator enables the analysis of time-domain signals in a higher dimensional space, making it easier to
distinguish between two segments of signals that appear similar in the time domain. In addition, periodic signals have distinctive characteristics in their Fourier transforms due to their discrete spectra. We provide a visualization of the DTFT of the state sequences in Appendix E, which reveals that the DTFT of the periodic state sequence is approximately discrete. This observation suggests that the periodic information of the signal can be explicitly extracted in the frequency domain, particularly for the periodic cases provided by Theorem 2. For non-periodic sequences, some regularity information can still be obtained by the frequency range that the signals carry.
In addition to those advantages, the operation of the Fourier transforms also yields a concise auxiliary objective similar to the TD-error loss [18], which we will discuss in detail in the following section.
5.2 Learning Objective of SPF
In this part, we propose our method, State Sequences Prediction via Fourier Transform (SPF), and describe how to utilize the frequency pattern of state sequences to learn an expressive representation. Specifically, our method performs an auxiliary self-supervision task by predicting the discrete-time Fourier transform (DTFT) of infinite-step state sequences to capture the structural information in the state sequences for representation learning, hence improving upon the sample efficiency of learning.
Now we model the auxiliary self-supervision task. Given the current observation  $s_t$  and the current action  $a_t$ , we define the expectation of future state sequence  $\widetilde{s}_t$  over infinite horizon as
$$
\left[ \widetilde {s} _ {t} \right] _ {n} = \left[ \widetilde {s} \left(s _ {t}, a _ {t}\right) \right] _ {n} = \left\{ \begin{array}{l l} \gamma^ {n} E _ {\pi , p} [ s _ {t + n + 1} | s _ {t}, a _ {t} ] & n = 0, 1, 2, \dots \\ 0 & n = - 1, - 2, - 3 \dots \end{array} . \right. \tag {5}
$$