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.
Paper
Similar papers
Peer review
Summary
The authors in this paper investigated the structural information in state sequences for reinforcement learning, and proposed the prediction of Fourier transformation for the input sequences, which is further used an auxiliary task for learning policies. Some theoritical results are provided to show the existence of structual information. Experimental results on MuJoKo tasks demonstrated the promise of the proposed method.
Strengths
1. The idea of performing prediction on the Fourier spectrum is interesting. This approach may better extract information hidden in the time domain, which constitutes the main contribution of this paper. 2. From the experimental results, the proposed auxiliary task indeed helped learn better policies for RL agents, w.r.t. both SAC and PPO.
Weaknesses
1. There seems to be a gap between the motivation and the implementation. The introduction and some theoretical results, e.g., Theorem 3, imply that the prediciton is on the Fourier domain; however, the algorithm (SPF) simply used a parameterized prediction module. Consequently, it's unclear to me whether the proposed algorithm can extract the information in the Fourier spectrum. 2. It's unclear to me whether the proposed method can capture the long-term information. From Figure 2, it only considers two adjacent states as inputs. 3. The significance in Theorem 1 may not be sufficient, as the upper bound can be loose. Note that the use of l1 norm for the difference between two distributions will make the term in Equation (3) much greater than 1. Furthermore, it's unclear whether the states distributions for two policies can vary that much, given the use of Rmax / (1 - \gamma).
Questions
1. Some notations in Figure 2 are not defined (e.g., $P$) and cannot be located from the definition of the prediction loss (e.g., $\mathcal{F}$) between L261 and L262. 2. What does $\Gamma_{Re}$ do in the prediction loss, given that $\mathcal{F}_{Re}$ is already real?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
2 fair
Presentation
2 fair
Contribution
2 fair
Limitations
Yes, the authors mentioned that the current version may not handle the image inputs.
Summary
Data-efficient RL exploits representation learning which learns representations by predicting long-term future states. However, inherent structural information in sequential state signals is ignored in existing methods. The authors propose SPF which exploits the frequency domain of state sequences. The proposed method performs an auxiliary self-supervision task to improve the efficiency of representation learning. This method outperforms several state-of-the-art algorithms in terms of both sample efficiency and performance.
Strengths
* Notations are clear and methods are clearly proposed.
Weaknesses
* Results are not promising compared to SAC-OFE. In Humanoid, SAC-OFE outperforms SAC-SPF, and in HalfCheetah, Walker2d, Ant, Swimmer, SAC-OFE is comparable to SAC-SPF. * I could not find critical superior points to use Fourier transform. As the authors pointed out in the conclusion, experiments on visual-based RL settings are required to show the benefits of the Fourier transform. * The inherent structural information in a sequential state is not clearly discussed in the paper. **Minors** * Typo in Figure2, There is no predictor F in the figure.
Questions
* What is inherent structural information we can get from Fourier-transform? * The main concern is that results do not outperform the baseline. Also, in 6.3, the results of all the baselines and SPF are comparable. ==================================================================================================== The authors successfully handle my concerns. I raise my score to borderline accept.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
1 poor
Presentation
2 fair
Contribution
2 fair
Limitations
Yes
Summary
Unlike many prior model-based RL methods that predict one future state at a time, the authors present a Fourier-transform based approach that predicts future states over a fixed interval. They prove that a large difference in policy performance will cause a significant difference between the experienced states of a good and bad policy. This conclusion is important for utilizing frequency-space analysis of the states. Their approach uses a self-supervised learning objective where they attempt to fit an online encoder given the current states, and a target encoder given the next states. Note that both the online and target encoder, given a single state, will predict frequencies for the rest of the trajectory. This is in contrast to prior methods that will just compute the next state. The TD-like error between predictions is then used to train the online encoder. I'm not sure I fully understand it, but it seems quite elegant in that their approach closely resembles the one-step TD error. As far as I can tell, they are not actually using the predicted state sequences for planning, as in model-based RL. Rather, they are learning the features necessary to predict future state sequences. These features are useful for a downstream policy. They evaluate their method on the standard MuJoCo/Gym benchmarks which outperforms baselines. They compare their method using SAC/PPO to SAC/PPO using OFE or no prediction at all. They continue to ablate nearly all portions of their approach.
Strengths
- The paper is very well written - Their approach is novel and elegant - Their experiments, specifically the ablation, is well-done - They show their approach outperforms baselines
Weaknesses
- They only compared to one other feature-learning algorithm. It would be beneficial to have additional baselines. - Unless I'm misunderstanding, it seems like they've purposely limited the applications of their approach to fully observable, model-free RL - They already predict future states, they might as do some online planning - This could be applied to partially observable domains to learn Markov states - They do not list any limitations of their approach
Questions
- "Therefore, state sequences maximally preserve the influence of the transition intrinsic to the environment and the effect of actions generated from the current policy" - I'd be careful saying this. In the case of a deterministic policy, yes. But if you add stochasticity (e.g. epsilon greedy or policy-gradient sampling) I'm not sure you can claim you can predict a future state sequence. You can perhaps predict the distribution of states. - "it is widely accepted that the frequency domain shows the regularity properties of the time-series data" - There are citations for this but I find it a bit hard to believe. For cyclical processes this makes sense, but there is no guarantee, for example, that you will see a specific feature every k timesteps. Could you add an intuitive explanation of why this should work for latent state sequences? - This is explained later in the paper, but perhaps it makes sense to say you show/prove this in the introduction - Eq. 5: Why do you discount the future state using gamma? I understand dicounting the reward/return, but not effectively "decaying" the future state. Could you briefly explain the reasoning here? Is it to weight nearby states more heavily?
Rating
8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
4 excellent
Presentation
4 excellent
Contribution
4 excellent
Limitations
- The authors do not list any limitations
Summary
The paper suggests a representation learning method that leverages the existing periodic structures in the state sequences of a policy. Specifically, they propose using the features obtained by the Fourier transform of the state sequence. To practically extract such features, they employ a Bellman equation-like recursive form of the Fourier transform (FT), which enables efficient online learning. The authors demonstrate that their representation learning method helps improve the performances of SAC and PPO in Gym MuJoCo environments.
Strengths
- The idea of using frequency features in sequential decision-making problems seems novel and sensible, given that real-world tasks often inherently have periodic (sub-)structures in their states (e.g., robotic locomotion). - The recursive objective in Equation (7), which enables learning Fourier features in practice without the need to deal with the entire state sequences, is intriguing and potentially inspirational for future research. - The paper includes an informative ablation study.
Weaknesses
- The Fourier features inherently depend on the policy and task, making them less reusable for multiple tasks, especially compared to previous self-supervised representation learning methods. - The theorems in the paper do not necessarily explain why Fourier features could be potentially useful for *representations*. Theorem 3 states a (very crude) bound for the performance difference between two policies expressed in the frequency domain. However, it is unclear as to why the features extracted in this manner are useful or sufficient for modeling policies/value functions as well, which is the way SPF utilizes the learned features. I would have expected a theorem that says the sub-optimality of using the learned Fourier representation decreases as the state sequence shows stronger periodicity.
Questions
- Can the authors elaborate on why Fourier features are useful for representations? For example, let's say a state sequence from a policy forms a perfect sinusoidal curve so that its Fourier transform collapses to a single point. In such a case, its Fourier representation (which is a constant) may not contain sufficient information for control (i.e., for producing optimal actions to maintain that sinusoidal curve). - Could SPF be harmful in tasks that do not have periodicity? Have the authors tested SPF on such environments/tasks?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The authors briefly mention the limitations of SPF in Section 7. However, it would be also important to include a discussion on the second question above in the limitation section as well.
Response to Authors
While I greatly appreciate the new results on Adroit, they have left me rather confused about the effectiveness of Fourier features. If there is no periodicity in trajectories, why would Fourier features be helpful? Perhaps an alternative hypothesis is that the performance gain comes *not* from Fourier features but from some practical implementation details, hyperparameters, etc., that simply lead to better performance or make the representations better. I also acknowledge that this benefit might come from detecting irregular, sudden changes (as the authors mentioned), but it would require more controlled analyses to prove this point (e.g., Do the learned Fourier features actually detect such sudden changes?, Why detecting such events is necessarily helpful for performance?, etc.). That being said, I wouldn't decrease my score as I believe the ideas and initial results are still valuable. Regarding the perfect sinusoidal curve example, my point is that since the Fourier representation $z$ is always a *constant* (i.e., always only non-zero at $\omega$), it might be insufficient for the *policy* $\pi(a|z)$. How can a policy produce optimal actions when the variable on which it is conditioned is always a constant?
Thanks for your comments
Thank you for your insightful comments and for helping us gain deeper insights into our work. We respond to your comments as follows and sincerely hope that our response could properly address your concerns. **Q1. About the effectiveness of Fourier features on Adroit tasks.** **A1.** Thank you for raising this point. By outputting the state sequences generated by our learned policy, we observed that the AdroitHandDoor task exhibits stronger periodicity compared to the AdroitHandPen task. This observation may help explain why our method performs more effectively in the Door task compared to the Pen task. **According to the properties of the Fourier transform, our method demonstrates better performance on tasks with strong periodicity and has the ability to capture sudden changes with periodicity.** **Q2: About the perfect sinusoidal curve example.** **A2.** We apologize for the misunderstanding. Here is our explanation. In the practical implementation of our method, we predict the Fourier transform (FT) of the future state sequence $\\{ s_{t+1},s_{t+2}, \cdots \\}$ based on the current state $ s_t$ and the current action $a_t$. Specifically, **for different states along the sinusoidal curves, the corresponding future state sequences start from different time steps, resulting in sinusoidal curves with varying time shifts, i.e., different phases.** Thus, in the case of perfect sinusoidal curves, the future state sequence predicted from state $s_t$ and action $a_t$ can be expressed as follows: $$ f(t)=\begin{cases} sin(\omega_0 t+\varphi)& \text{t}>0, \\ 0& \text{t}\leq 0 \end{cases} =sin(\omega_0 t+\varphi) H(t), $$ where $\omega_0$ means the frequency, $\varphi$ means the phase, and $H(t)$ means the Heaviside step function. The Heaviside step function takes on a value of zero for negative arguments and a value of one for positive arguments. According to the convolution theorem, the Fourier transform of $f(t)$ is given by: $$ F(\omega)=-\cos \varphi \cdot\frac{\omega_0}{\omega^2-\omega_0^2}+\frac{\pi}{2}\sin \varphi \left( \delta(\omega-\omega_0)+\delta(\omega+\omega_0) \right)-i\left[ \sin \varphi \cdot\frac{\omega}{\omega^2-\omega_0^2}+\frac{\pi}{2}\cos \varphi\left( \delta(\omega-\omega_0)+\delta(\omega+\omega_0) \right) \right], $$ where we use Fourier transform in the non-unitary form with angular frequency. Note that the phase $\varphi$ plays a role in determining the FT values of future state sequences that start from different time steps. The encoder and the predictor capture such phase information when predicting the FT of the future state sequences. Thus, **the representations of the states along the sinusoidal curves can be distinguished by their respective phases.** We agree that the states spaced one cycle apart share the same representations, which means the agent exhibits periodic patterns of behavior and performs the same actions selected by the policy $\pi(\cdot|\phi(s))$ when encountering these states with the same phase. **We recognize that the influence of the phase $\varphi$ may be relatively weaker compared to the dominant frequency, as influenced by the delta function.** Therefore, in practical implementations, we address this by either concatenating the original state with Fourier features or training the encoder using the critic loss and prediction loss simultaneously.
Thank you for the detailed clarification on the perfect sinusoidal curve example!
Thank you for your kind support
Thank you again for your valuable comments and constructive suggestions, which are of great help to improve the quality of our work. We sincerely hope that our rebuttal could properly address your concerns. If you have further concerns, please let us know, and we will continue actively responding to your comments and improving our submission.
Thank you for adding the additional baselines and fully addressing my concerns. I was unclear in my comment on partial observability and online planning, I apologize. My understanding is that your approach could allow predicting future states without sequentially rolling them out. For example, if I wanted to predict $s_{t + 100}$ using traditional MBRL, I would need to loop through my dynamics model 100 times, potentially accruing error along the way. Your current approach conditions the model on all future actions (Eq. 5), but if it instead conditioned on the current action and policy (i.e., $s_{t+L} = g(s_t, a_t, L | \pi)$, one could effectively predict $s_{t + 100}$ in one forward pass, right? Or perhaps I am misunderstanding. This seems like it could be a powerful property if I were training a Dreamer-style model with partial observability, as one could learn a latent state representation $s_t = f(o_1, \dots, o_t)$ that predicts future observations via $g(s_t, a_t, 100,000 | \pi) = o_{t+100,000} $, without worrying about computational cost. Additionally, one could do fast online planning of long trajectories that would not be possible with RNN or transformer-based dynamics models.
Thanks for your comments
Thank you for the insightful comment! **We believe that our method (SPF) can do fast online planning in a Dreamer-style algorithm, in which SPF could generate multi-step future states without using RNN or transformer-based dynamics models.** The reasons are as follows. + Our method could predict future states without sequentially rolling them out. However, **we would like to point out that our method conditions the prediction model on the current state $s_t$ and the current action $a_t$, rather than all future actions.** According to Equation (5) in our main text, we predict the expectation of the future state sequences $\\{s_{t+1}, s_{t+2}, \cdots\\}$ conditioned on the current state $s_t$ and the current action $a_t$. Note that the expectations in Equations (5) and (7) are computed with respect to the policy $\pi$ and the dynamic $p$, indicating that the predicted state sequences are also conditioned on the policy. + In practical implementation, the inputs of the predictor module are the representations of the current state $s_t$ and the current action $a_t$. The output of the predictor is an $L-$dimensional vector, which corresponds to $L$ discrete samples of the discrete-time Fourier transform (DTFT) of the future state sequences. According to the inverse discrete Fourier transform (IDFT), we can compute the future state sequences from the predicted DTFT, thereby enabling the estimation of $s\_{t+L}$ for $L=1,2,\cdots,H$ in one forward pass. **In other words, by leveraging the IDFT, the Fourier predictor can generate multi-step future states for fast online planning.**
Thanks to the authors for the response, I really appreciate it. Regarding the gap between motivation and implementation, it is still unclear to me how well the module $\mathcal{F}$ can approximate the real Fourier transformation. For example, how about their distance between prediction and ground truth? For the upper bound in Theorem 1, my point is that the existence of $\frac{R_max}{1 - \gamma}$ just weakens the claim, as it could dominate the performance gap between $\pi_1$ and $\pi_2$, which makes the last term less relevant.
Thanks for your comments.
Thank you for your insightful comments and for helping us gain deeper insights into our work. We respond to your comments as follows and sincerely hope that our response could properly address your concerns. **Q1. How well the module $\mathcal{F}$ can approximate the real Fourier transformation?** **A1.** + Given that 1) the parameters of the policy $\pi$ are updated alternately during the training of the prediction module $\mathcal{F}$; 2) the policy $\pi$ is trained to approach the optimal policy $\pi^*$. The ground truth for our prediction module $\mathcal{F}$ is the real Fourier transform (FT) of the state sequences generated by the optimal policy $\pi^*$. **According to the last inequality of the proof provided in the global response, the distance between the predicted FT and the ground truth is constrained by the objective of our method (SPF).** This can be expressed as follows: $$ \begin{align} \| F_{\pi_*} - F_{\overline{\pi}\circ\phi}\| \leq (1-\gamma)^{-1}\cdot E_{(s,a,s')\sim\mathcal{D}}\left[s'+e^{-j\omega}\gamma E_{a'\sim{\overline{\pi}\circ\phi}}\left[F_{\overline{\pi}\circ\phi}(s',a')\right] - F_{\overline{\pi}\circ\phi}(s,a)\right]. \end{align} $$ Here, $F_{\pi_*}$ represents the real FT of the state sequence generated by the optimal policy $\pi_*$, and $F_{\overline{\pi}\circ\phi}$ represents the output of the prediction module $\mathcal{F}$, which is influenced by the current policy $\overline{\pi}$ and the current encoder $\phi $. According to the inequality we mentioned above, the distance between the predicted FT of our module $\mathcal{F}$ and the real FT is minimized when we minimize our objective $\overline{\delta}\_\pi(s,a,s'):=s'+e^{-j\omega}\gamma E\_{a'\sim{\pi}}\left[F\_{\pi}(s',a')\right] - F\_{\pi}(s,a)$ and collect a sufficient number of samples in the replay buffer $\mathcal{D}$. + **We present a comparison table between the real Fourier transform (FT) and the predicted FT.** To obtain these results, we use our trained policies to interact with six MuJoCo environments and record the states and the actions for the 201 last steps of an episode, denoted as $\\{s_t\\}\_{t=0}^{200}$ and $\\{a_t\\}\_{t=0}^{200}$, respectively. **The real FT** of the state sequence $\\{s_1,s_2,\dots,s_{200}\\}$ over a 200-step horizon is computed using the formulation of the Discrete-time Fourier transform (DTFT). For the predicted FT, we employ our trained encoder and prediction module, with the state $s_0$ and the action $a_0$ serving as inputs. The output of the prediction module represents **the predicted FT** of the state sequence $\\{s_1,s_2,\dots,s_{200}\\}$. We compute the cosine similarity between $A$ and $B$ with the formula $-\frac{A\cdot B}{\\|A\\|\\,\\|B\\|}$, where values closer to $-1$ indicate greater similarity while values closer to $1$ indicate greater dissimilarity. The cosine similarity between the real FT and the predicted FT is listed below. | Environments | Cosine similarity between the real FT and the predicted FT | | ------------ | ---------------------------------------------------------- | | HalfCheetah | -0.655 | | Hopper | -0.660 | | Walker2d | -0.651 | | Ant | -0.351 | | Swimmer | -0.557 | | Humanoid | -0.651 | **Q2. The existence of $\frac{R_\text{max}}{1-\gamma}$ in Theorem 1.** **A2.** Thank you for raising this point! **We acknowledge that when $\frac{R_\text{max}}{1-\gamma}$ takes on a large value, our claim that state sequences exhibit significant differences in the presence of a significant difference in policy performance is weakened.** An alternative viewpoint regarding Theorem 1 is that as the distribution of state sequences generated by the trained latent policy $\overline{\pi}\circ\phi$ approaches the distribution of state sequences generated by the optimal policy $\pi^*$, the performance of our trained policy also approaches that of the optimal policy.
Decision
Accept (spotlight)