Skip to content

Forecasting in Offline Reinforcement Learning for Non-stationary Environments

🕒 Published (v1): 2025-12-01 18:45 UTC · Source: Arxiv · Venue: NEURIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

FORL (Forecasting in Non-stationary Offline RL) addresses the failure of offline RL policies when deployed in environments with unknown, time-varying additive observation offsets. It combines a conditional diffusion model (for multimodal in-episode state belief updates) with a zero-shot time-series foundation model (Lag-Llama) to forecast and compensate for unseen offsets at test time, requiring no policy retraining or environment interaction.

Problem

Offline RL policies trained on stationary datasets degrade sharply when deployed in environments where the observation function is corrupted by unknown, episode-level additive offsets that evolve as a non-Markovian time series. Existing robust offline RL methods (e.g., DMBP, RORL) handle test-time noise or adversarial perturbations but assume Gaussian/parametric perturbation models and do not account for persistent, cross-episode, real-world non-stationarity in the observation space.

Method

FORL operates entirely at test time with a fixed offline-trained policy and two components:

  1. FORL Diffusion Model (DM): A conditional diffusion model trained on stationary offline data to model the multimodal distribution \(p(s_t \mid \tau^{(t,w)})\), where \(\tau^{(t,w)}\) is a window of (delta-observation, action) pairs. These delta-observations are offset-invariant, allowing the model to generate \(k\) candidate states reflecting where the agent plausibly is given its interaction history.

  2. Zero-Shot FM (Lag-Llama): A probabilistic univariate time-series foundation model used zero-shot to forecast the next \(P\) episode offsets from \(C\) past revealed offsets, producing \(l\) samples per episode per state dimension.

  3. Dimension-wise Closest Match (DCM): A fusion heuristic that, for each state dimension \(d\), selects the forecaster sample \(y_{j^*(d)}\) whose value is closest (in that dimension) to any diffusion candidate. This approximates the product distribution of the two sources without parametric assumptions. For the first \(w\) steps (before enough history), only the forecaster mean is used.

Key Contributions

  • Framework unifying zero-shot time-series forecasting with conditional diffusion-based state belief for offline RL under non-Markovian, episodic observation offsets.
  • Novel benchmark: five real-world time-series datasets (from distinct domains) injected as per-dimension offsets into D4RL and OGBench environments.
  • DCM fusion strategy: non-parametric, hyperparameter-free, and shown empirically to produce lower maximum prediction errors than KDE, mean/median selection, or max-likelihood variants.
  • Policy-agnostic plug-and-play: demonstrated across DQL, TD3+BC, RORL, IQL, FQL without retraining the FORL components.
  • Ablation showing that FORL-DM alone (no historical offsets needed) achieves 151.4% relative improvement over DQL baseline, and the full FORL yields a further 97.8% gain.

Results

  • maze2d-medium: FORL avg 74.1 vs. DQL 10.0, DQL+LAG-sÌ„ 37.2, DMBP+LAG 34.6.
  • maze2d-large: FORL avg 38.8 vs. DQL 3.5, DQL+LAG-sÌ„ 7.1, DMBP+LAG 12.1.
  • antmaze-umaze-diverse: FORL avg 75.1 vs. DQL 16.1, DQL+LAG-sÌ„ 44.9, DMBP+LAG 51.4.
  • antmaze-large-diverse: FORL avg 34.6 vs. DQL 6.0, DQL+LAG-sÌ„ 12.0, DMBP+LAG 10.5.
  • cube-single-play (FQL base): FORL-F avg 45.7 vs. FQL 0.1, FQL+LAG-sÌ„ 5.2.
  • antmaze-large-navigate: FORL avg 42.2 vs. FQL 18.1, FQL+LAG-sÌ„ 16.7.
  • Policy-agnostic (maze2d-large, RORL base): FORL-R avg 71.9 vs. RORL 2.8, RORL+LAG-sÌ„ 21.6.
  • DCM maximum prediction error: 2.40 vs. FORL(MAX) 9.33.
  • FORL outperforms all baselines across all 5 real-world non-stationarity domains consistently (Welch's t-test, p < 0.05).

Limitations

  • Assumes additive observation perturbations only; more general observation transformations (multiplicative, compositional) are not handled.
  • Requires past ground-truth offsets to be revealed after each episode; when this is unavailable, the H-LAG+DCM variant degrades (DM-only is more stable but lower peak performance).
  • Zero-shot forecaster accuracy (Lag-Llama) is a bottleneck; when real-world offsets deviate sharply from the forecaster's prior, performance degrades.
  • Offset assumed constant within each episode; intra-episode non-stationarity is addressed but is a harder setting with reduced performance.
  • Evaluated only on continuous locomotion/navigation/manipulation benchmarks; applicability to language-grounded or discrete-action agents is unexplored.

Relevance to Agentic AI / LLM Agents

FORL is directly relevant to the challenge of deploying trained agents in distribution-shifted real-world environments without retraining — a core concern for any agentic system operating outside its training distribution. The use of a zero-shot time-series foundation model (Lag-Llama) as a drop-in forecasting module is an instance of the broader trend of composing specialist foundation models to extend agent robustness without task-specific fine-tuning. The DCM fusion strategy offers a lightweight recipe for combining multimodal belief estimates with unimodal priors, a pattern transferable to agent architectures that must reconcile multiple uncertain information sources. For LLM-based agents that may receive corrupted or shifted tool outputs across sessions, the episodic partial-observability framing here provides a formal grounding for how to model and compensate for such drift.