Wavelet Predictive Representations for Non-Stationary Reinforcement Learning

TL;DR

This work addresses the challenge of non-stationary reinforcement learning by modeling how tasks evolve within a sequence of MDPs. It introduces WISDOM, a framework that learns a trainable wavelet representation network $Y_{\phi}$ to transform time-varying task representations $\mathbf{z}$ into the wavelet domain, capturing multi-scale trends and rapid variations. A wavelet TD update ensures stable tracking of task evolution, and this representation is integrated into a SAC-based policy, with theoretical results guaranteeing contraction and policy improvement. Empirical results on Meta-World, MuJoCo, and Type-1 Diabetes demonstrate superior sample efficiency and robust performance under irregular non-stationarities, highlighting the practical value of wavelet-based task representations in NSRL.

Abstract

The real world is inherently non-stationary, with ever-changing factors, such as weather conditions and traffic flows, making it challenging for agents to adapt to varying environmental dynamics. Non-Stationary Reinforcement Learning (NSRL) addresses this challenge by training agents to adapt rapidly to sequences of distinct Markov Decision Processes (MDPs). However, existing NSRL approaches often focus on tasks with regularly evolving patterns, leading to limited adaptability in highly dynamic settings. Inspired by the success of Wavelet analysis in time series modeling, specifically its ability to capture signal trends at multiple scales, we propose WISDOM to leverage wavelet-domain predictive task representations to enhance NSRL. WISDOM captures these multi-scale features in evolving MDP sequences by transforming task representation sequences into the wavelet domain, where wavelet coefficients represent both global trends and fine-grained variations of non-stationary changes. In addition to the auto-regressive modeling commonly employed in time series forecasting, we devise a wavelet temporal difference (TD) update operator to enhance tracking and prediction of MDP evolution. We theoretically prove the convergence of this operator and demonstrate policy improvement with wavelet task representations. Experiments on diverse benchmarks show that WISDOM significantly outperforms existing baselines in both sample efficiency and asymptotic performance, demonstrating its remarkable adaptability in complex environments characterized by non-stationary and stochastically evolving tasks.

Figures (15)

• Figure 1: A motivating example.
• Figure 2: The architecture of WISDOM. Module A begins with task inference to derive time-domain task representation $z$. Then module B transforms $z$ into wavelet domain by a wavelet representation network $Y_\phi$, jointly optimized by wavelet TD loss and AR loss to derive wavelet task representation $\hat{z}$. Finally, module C integrates $\hat{z}$ to adjust the policy based on predicted evolving trend.
• Figure 3: Testing average performance on Meta-World over 6 random seeds. Our WISDOM achieves rapid convergence and exhibits excellent asymptotic performance.
• Figure 4: Testing average return on glucose control environments over 6 random seeds.
• Figure 5: Testing return on MuJoCo over 6 seeds.

Theorems & Definitions (12)

• Theorem 1
• Theorem 2
• Theorem 3
• Lemma 1
• Proof 1
• Theorem 1
• Proof 2
• Lemma 2
• Theorem 2
• Proof 3