Forecasting in Offline Reinforcement Learning for Non-stationary Environments

AI summary

Overview

Research area: Offline reinforcement learning, robust decision-making under distribution shift, time-series forecasting with foundation models, and diffusion-based generative modeling.

Technical level: Advanced. The paper assumes familiarity with MDPs and POMDPs, offline RL datasets (D4RL, OGBench), denoising diffusion probabilistic models, and zero-shot time-series foundation models.

Scope: The paper proposes Forl (Forecasting in Non-stationary Offline RL), a test-time framework that combines a zero-shot time-series forecaster with a conditional diffusion model to recover an agent's true state when its observations are corrupted by unknown, per-episode additive offsets.

What This Paper Is About

Offline RL agents are trained on a fixed, pre-collected dataset and are expected to generalize at deployment. In practice, the mapping from true state to observation can shift over time — an industrial robot may receive a new calibration offset each day, or a sensor may drift until the next recalibration — so the agent perceives a displaced version of where it actually is. Existing offline RL methods either assume the environment is stationary or handle only synthetic test-time perturbations, and they cannot anticipate abrupt, non-Markovian shifts that follow real-world time-series patterns.

The goal of Forl is to let an agent that was trained on clean, fully observed, stationary data estimate its true state from the very first step of each new episode at test time, using only (a) a short history of previously revealed offsets and (b) the within-episode record of its own actions and observed state changes, without any environment interaction, retraining, or policy updates.

Key Contributions

1. A non-stationarity-agnostic framework. Forl accommodates future observation offsets without assuming any specific pattern of non-stationarity during training, so no retraining or hyperparameter tuning is needed when the agent meets unseen non-stationary patterns at test time. It targets non-trivial non-stationarity without requiring environment interaction or knowledge of POMDPs during training.

2. A modular architecture combining forecasting and belief generation. The framework unifies a zero-shot time-series forecasting foundation model (Lag-Llama) that predicts future offsets from a limited context of past offsets, with a conditional diffusion model (Forl-Dm) trained on stationary offline data that generates a multimodal set of plausible candidate states from the in-episode sequence of action and delta-observation pairs.

3. Dimension-wise Closest Match (DCM). A lightweight, non-parametric fusion strategy that selects, for each state dimension, the forecast sample whose value is closest to any sample produced by the diffusion model, blending the multimodal diffusion belief with the unimodal forecast.

4. A new benchmark. The paper introduces a benchmark that injects offsets derived from real-world time-series datasets into standard offline RL benchmarks (D4RL and OGBench), and reports that Forl consistently outperforms competitive baselines across it.

Main Findings

• Forl improves average normalized scores across most benchmark suites. Averaged over the five real-data domains, Forl reaches 74.1 on maze2d-medium versus 37.2 for DQL+Lag-s̄, 34.6 for Dmbp+Lag, and 10.0 for plain DQL; 38.8 on maze2d-large versus 7.1, 12.1, and 3.5; 75.1 on antmaze-umaze-diverse versus 44.9, 51.4, and 16.1; 48.7 on antmaze-medium-diverse versus 42.1, 40.1, and 18.9; 34.6 on antmaze-large-diverse versus 12.0, 10.5, and 6.0; and 21.3 on kitchen-complete versus 18.1, 16.7, and 13.8.

• Forl-F likewise leads on the flow-based policy settings. On cube-single-play, Forl-F averages 45.7 versus 5.2 for FQL+Lag-s̄ and 0.1 for plain FQL; on antmaze-large-navigate it averages 42.2 versus 21.2 and 12.0.

• The forecasting-plus-robust-RL pipeline alone is weak. Dmbp+Lag, which first predicts offsets with a time-series model and then applies a noise-robust offline RL algorithm, is consistently behind Forl, indicating that refining forecasted states is not sufficient when the offsets are abrupt and non-Markovian.

• The diffusion component carries much of the benefit even without past offsets. When no historical offsets are available, Forl-Dm performs on par with the extended baselines that do use historical offsets and forecasting, and achieves a 151.4% improvement over DQL. The full Forl framework yields a 97.8% relative performance improvement over Dm alone.

• DCM is the strongest fusion rule tested. Compared with Forl(KDE), DM-FS-s̄, DM-FS-š, and Forl(max), all of which combine the same two sample sets, DCM attains the lowest maximum prediction error (2.40) versus Forl(max) at 9.33, and the most stable mean error, as measured by the ℓ2 norm between ground-truth states and predicted states.

• Historical-offset variants reveal a trade-off. Feeding Dm-generated offset history into the forecaster (H-Lag) is improved by adding DCM (H-Lag + dcm), but overall scores are better on average when using Dm samples alone, while H-Lag + dcm is more stable.

• The recovery mechanism is interpretable. In the maze illustration, after moving north several steps and hitting a wall, the possible ground-truth states reduce to the regions three steps south of walls — four candidate regions — which the diffusion model captures, whereas DQL+Lag-s̄ places its estimate significantly farther from the actual state.

Methodology in Plain English

Training. The agent is trained entirely offline on a stationary dataset with fully observable states, such as D4RL. Two models are learned from that same data: an off-the-shelf offline RL policy (Diffusion-QL, or Flow Q-learning in the Forl-F variant) and a conditional diffusion model. The diffusion model is trained to denoise true states conditioned on a short window of past (delta observation, action) pairs, where the delta observation is the difference between consecutive observations. Because the offsets are additive and indistinguishable from state differences within an episode, these deltas equal the corresponding true state changes.

The core intuition. An offset is constant within an episode but unknown, so the agent cannot separate its true state from the bias from a single observation. However, the sequence of actions it took and the changes it observed constrains where it could be. If the agent moved north for three steps and then hit a wall, its true position must be just south of a wall. This "retrospection" is what the diffusion model supplies: a multimodal set of candidate states consistent with the observed history.

Test time. Each episode, the forecaster (Lag-Llama, a probabilistic zero-shot model applied independently per dimension) predicts the offset for the next P episodes from the last C observed offsets; l samples are drawn per episode. For the first w steps, before enough history exists, the state estimate is simply the observation minus the mean forecast offset. Once w steps have elapsed, the diffusion model generates k candidate states from the windowed history, the forecaster yields l candidate states (observation minus each forecast offset), and DCM fuses them dimension by dimension: for each dimension it picks the forecast value whose closest diffusion sample is nearest. That fused vector is the state fed to the policy. The process repeats within the episode, with true offsets never revealed during evaluation.

Why This Matters

Impact on research. The paper reframes test-time robustness in offline RL from "reject noise" to "forecast and reason about a slowly revealed latent bias," and it introduces the observation function as the non-stationary component to study, a scope the authors describe as underexplored. It also supplies a concrete benchmark that couples real-world time-series data with standard offline RL datasets, giving the community a reproducible way to measure episodic, non-Markovian observation shift rather than Gaussian noise.

Real-world applications.

• Industrial robotics, where joints receive daily calibration offsets that stay fixed within an operating interval.
• Sensor networks, where a device exhibits a deviation until its next scheduled recalibration.
• Healthcare and finance, where data may be partially aggregated or withheld for regulatory compliance, leaving a single representative offset as the dominant per-episode factor.
• Bandwidth-constrained or privacy-sensitive deployments, where storing only episodic calibration values instead of continuous interaction buffers is preferable to full replay storage.

Industry relevance. The approach requires no retraining, no hyperparameter tuning, and no online exploration, which matters for safety-critical systems where trial-and-error is unacceptable. Storing only per-episode offsets rather than continuous buffers reduces data-handling burden and fits privacy-constrained settings. Because Forl is policy-agnostic and plugs into existing offline RL algorithms, it can be layered onto already-trained diffusion-based or flow-based policies.

Future Directions

• Operating with no past offsets at all. The paper evaluates variants for this case (Forl-Dm alone, H-Lag, H-Lag + dcm, Med + dcm, Med + Noise); the open question is how to close the remaining gap between Dm's better average scores and H-Lag + dcm's greater stability.
• Intra-episode non-stationarity. The paper lists this as an explicit experimental question; extending the framework beyond constant-per-episode offsets to offsets that change within an episode is unresolved.
• Multivariate forecasting. Lag-Llama is univariate and the authors forecast each offset dimension independently, leaving open whether multivariate time-series foundation models would capture cross-dimension structure and improve accuracy.
• Broader plug-and-play validation. Forl is demonstrated with Diffusion-QL and Flow Q-learning; whether the same framework transfers unchanged to other offline RL policy classes without retraining remains to be tested.

Target Audience

Researchers and practitioners working on offline RL, robust RL under distribution shift, and partially observable decision-making; engineers deploying learned policies in robotics or industrial sensing where calibration drift and sensor bias are routine; and time-series researchers interested in how zero-shot forecasting foundation models can be coupled with generative state estimation. Readers without a background in diffusion models or POMDPs will find the method sections demanding, though the maze illustrations and the overview figure convey the central idea accessibly.

Authors’ abstract

Offline Reinforcement Learning (RL) provides a promising avenue for training policies from pre-collected datasets when gathering additional interaction data is infeasible. However, existing offline RL methods often assume stationarity or only consider synthetic perturbations at test time, assumptions that often fail in real-world scenarios characterized by abrupt, time-varying offsets. These offsets can lead to partial observability, causing agents to misperceive their true state and degrade performance. To overcome this challenge, we introduce Forecasting in Non-stationary Offline RL (FORL), a framework that unifies (i) conditional diffusion-based candidate state generation, trained without presupposing any specific pattern of future non-stationarity, and (ii) zero-shot time-series foundation models. FORL targets environments prone to unexpected, potentially non-Markovian offsets, requiring robust agent performance from the onset of each episode. Empirical evaluations on offline RL benchmarks, augmented with real-world time-series data to simulate realistic non-stationarity, demonstrate that FORL consistently improves performance compared to competitive baselines. By integrating zero-shot forecasting with the agent's experience, we aim to bridge the gap between offline RL and the complexities of real-world, non-stationary environments.