🤖 machine learning

Dynamics Models for Offline Hyperparameter Selection in Real-World RL

This paper presents the first real-world application of offline calibration models for hyperparameter selection in reinforcement learning, demonstrating their effectiveness in a municipal water treatment plant by generating realistic long-horizon rollouts and recovering meaningful sensitivity trends on high-dimensional, non-stationary sensor data.

Original authors: Jordan Coblin, Han Wang, Martha White, Adam White

Published 2026-08-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jordan Coblin, Han Wang, Martha White, Adam White

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to do a complex job, like driving a car or managing a power plant. You can't just guess the settings for its brain; you have to tune them perfectly. This is called "hyperparameter selection." Usually, scientists do this by letting the robot practice in a video game simulator. But what if the real world is too dangerous, expensive, or slow to simulate? What if you can't build a perfect video game of a water treatment plant because the physics are too messy?

This is where a clever trick called "calibration models" comes in. Instead of building a fake world from scratch, these models act like a time-traveling mirror. They look at a massive library of old recordings (offline data) from the real system and try to predict what would happen next if a robot took a specific action. The goal isn't to be a perfect crystal ball that predicts the future with 100% accuracy; the goal is just to be a good enough judge. If the model can tell you which robot settings work best in its "mirror world," those same settings should work well when you finally send the robot out into the messy, real world.

The Paper: A Mirror for Water Plants

In this paper, a team of researchers decided to test this "mirror" idea in the real world for the very first time. They didn't use a simple video game; they went to a municipal water treatment plant in Drayton Valley, Alberta. This place is a chaotic mix of high-dimensional sensor data, changing weather, and non-stop noise. The researchers wanted to see if they could use a "calibration model" to help pick the best settings for an AI agent that predicts what the plant's sensors will do next.

The Challenge: The "Nexting" Game
The task wasn't to control the water (like turning valves on and off). Instead, the AI had to play a prediction game called "nexting." Imagine you are watching a river level rise. The AI's job is to look at the current data and guess the entire future path of that river level for a long time ahead. It's like trying to guess the plot of a movie just by looking at the first few minutes. The AI needs to get this right so that if it were eventually used to control the plant, it would know what to expect.

The Experiment: Building the Mirror
The team built four different types of "mirrors" (models) to simulate the water plant's behavior using a year's worth of sensor logs.

The K-Nearest Neighbors (KNN) Models: These are like a librarian who finds the most similar past days in the history books. If the plant is acting weird right now, the model looks for days in the past that looked exactly the same and says, "Okay, on those days, the pressure went up like this." They tried two versions: one that just looked at raw numbers (Euclidean) and one that looked at the deeper structure of the data (Laplacian).
The Neural Networks: These are the standard, complex AI brains (a Feedforward NN and a Gated Recurrent Unit) that try to learn the rules of the water plant from scratch.

They tested these models by rolling them forward in time for 30,000 steps (a long time in computer years) to see if they could generate realistic sensor readings without falling apart.

What They Found
The results were a mix of "promising" and "it's complicated."

The KNN Librarians Won the Long Race: The neural networks (the complex AI brains) started to collapse and give nonsense answers after about 50 to 100 steps. They couldn't handle the long-term prediction. However, the KNN models, especially the Laplacian version, were much more stable. They could generate long, realistic-looking trajectories that looked very similar to the real sensor data.
The Mirror Worked for Tuning: The most important test was whether the mirror could help pick the right "learning rate" (how fast the AI learns). The researchers found that the KNN models could successfully identify which learning rates were good and which were bad, matching the trends seen in the real data. This suggests that even if the model isn't perfect, it's good enough to tell you which settings to use.
The "Big Data" Question: They also tested if using a whole year of data (3.2 million samples) instead of just one week made the model better. The results were mixed. The big model could capture a wider variety of weird weather events, but it didn't consistently outperform the smaller model in every single test. It suggests that having more data helps, but it's not a magic bullet.
The "Time Travel" Problem: The researchers tried to see if the model could handle "distribution shift"—basically, what happens if the plant changes because of a new season or a broken sensor? They found that the model struggled to perfectly simulate a future that was very different from its training data. While it could mimic general changes, it couldn't perfectly predict specific, unique shifts without help.

The Bottom Line
This paper is a "proof of concept." It shows that for the first time, we can use these offline "mirror" models in a real, messy industrial setting to help tune AI agents. The KNN approach, which relies on finding similar past moments rather than learning complex rules, proved to be surprisingly robust for long-term predictions.

However, the authors are careful not to call this a solved problem. They point out that while the models can preserve the ranking of good settings (telling you "Option A is better than Option B"), they still struggle to perfectly simulate the future if the world changes in unexpected ways. It's a solid step forward, showing that we can use old data to prepare AI for the real world, but there is still work to be done to make these mirrors reflect every possible future perfectly.

Technical Summary: Dynamics Models for Offline Hyperparameter Selection in Real-World RL

Problem Statement

Deploying reinforcement learning (RL) in real-world industrial systems is frequently hindered by the difficulty of hyperparameter selection. Standard approaches face significant limitations: default hyperparameters rarely match specific system dynamics; simulator-based tuning requires accurate models that are often unavailable; and direct online tuning is often infeasible due to the high cost, risk, or time intensity of real-world interactions.

While "calibration models"—dynamics models trained on offline data to serve as surrogate environments for hyperparameter selection—have been proposed (Wang et al., 2022), they have previously been evaluated only in simple simulated domains. It remains unclear whether these models can effectively approximate complex, non-stationary, real-world dynamics to preserve the relative ranking of hyperparameter configurations. This paper addresses the gap between theoretical calibration models and their application in a real-world industrial setting.

Methodology

Framework: Data2Online

The study utilizes the "Data2Online" framework (Wang et al., 2022). In this setting, a calibration model $\hat{p}$ is trained on offline logs to approximate the true environment dynamics $p$ . The model is not used to optimize a policy directly but serves as an evaluation environment to select hyperparameters ( $\lambda$ ) for an agent that will subsequently learn online. The objective is to ensure that the optimal hyperparameters found in the model ( $\lambda^*_{model}$ ) match those of the true environment ( $\lambda^*_{env}$ ).

Task: Nexting Prediction

The specific task evaluated is nexting prediction, a General Value Function (GVF) problem (Modayil et al., 2012). Instead of controlling the plant, a passive agent predicts the discounted sum of future values for specific sensor signals (cumulants). Performance is measured by the Root Mean Squared Error (RMSE) between predicted values and empirical Monte Carlo returns, rather than by control return.

Calibration Model Architectures

The authors evaluate several model architectures:

k-Nearest Neighbors (kNN): A non-parametric approach that avoids extrapolation by predicting next states and rewards using only transitions within the dataset. Two variants are tested:
- Euclidean kNN: Uses raw state-action distance.
- Laplacian kNN: Uses an approximate Laplacian representation (Wu et al., 2018) to construct a distance metric sensitive to the underlying MDP structure.
- Ensembling: To address epistemic uncertainty, a Leave-One-Block-Out (LOBO) ensemble of five kNN models is used. Performance is aggregated by taking the worst rank across models to provide a conservative estimate.
Parametric Baselines: A Feedforward Neural Network (NN) and a Gated Recurrent Unit (GRU) trained to minimize one-step next-state prediction loss.

Evaluation Metrics

Rollout Quality: Qualitative visual comparison of long-horizon (30k steps) trajectories against true data.
Hyperparameter Sensitivity: Comparison of learning rate sensitivity curves (NRMSE vs. learning rate) between the calibration model and the "Online" setting (replayed held-out data).
Dynamic Time Warping (DTW): Used to measure similarity between model-generated and true trajectories, accounting for differences in speed, timing, or phase.

Experimental Setup

The study applies these methods to a municipal water treatment plant (WTP) in Drayton Valley, Alberta.

Dataset: Over two years of sensor logs (2022–2024) at 1 Hz, reduced to a 142-dimensional feature vector.
Sensors: Three critical sensors were selected: Membrane Pressure (PIT300), Influent Temperature (TIT101), and Influent Turbidity (TUIT101).
Scenarios:
1. Standard: Training on a one-week slice (~350k transitions).
2. Scaling: Training on a full year of data (~3.2M samples after subsampling).
3. Fine-Tuning: Pre-training an agent on 6 months of data and using the model to select fine-tuning learning rates.
4. Distribution Shift: Testing on data from April, May, and July 2023 to evaluate generalization to seasonal changes and sensor drift.

Key Results

1. Real-World Application and Rollout Quality

kNN Superiority: The kNN models (both Euclidean and Laplacian) successfully generated realistic long-horizon rollouts for the PIT300 sensor, closely resembling true data. In contrast, parametric baselines (NN and GRU) tended to collapse after 50–100 timesteps.
Laplacian vs. Euclidean: The Laplacian kNN showed closer overall alignment for PIT300, though the Euclidean kNN aligned more closely with Online curves for TIT101 and TUIT101. This suggests that the Laplacian representation may require tuning specific to the sensor dynamics.

2. Hyperparameter Sensitivity

Rank Preservation: Both kNN models generally recovered the correct hyperparameter rankings (learning rate sensitivity) compared to the Online setting.
Parametric Failure: The NN and GRU models failed to preserve hyperparameter sensitivity, consistent with their poor rollout quality.
Conclusion: Calibration models can support offline selection of learning rates for TD(0) prediction agents.

3. Scaling and Distribution Shift

Data Volume: Scaling the kNN model to a year-long dataset (12-month kNN) allowed it to capture a broader range of dynamics compared to the 1-week model. However, quantitative results using DTW were mixed; the 12-month model performed best on TIT101, while the 1-week model performed better on TUIT101.
Fine-Tuning: In a fine-tuning scenario, the 6-month kNN model was able to reflect a generic distribution shift, showing that performance of the smallest learning rates deteriorates as the deployment period moves further away, indicating a need for more adaptation. However, the model struggled to capture specific shifts, showing a mismatch with the Online sensitivity curves for later deployment periods.
Start State Selection: Strategies for selecting rollout start states (e.g., selecting nearest neighbors from the deployment period) showed limited benefit over random selection in this specific setup, suggesting that targeted start state selection alone is insufficient to fully simulate arbitrary distribution shifts.

Significance and Claims

The paper presents the first application of calibration models for offline hyperparameter selection in a real-world industrial setting. Its primary significance lies in providing a proof of concept that offline dynamics models can support RL deployment in environments where simulators are unavailable and online experimentation is costly.

The authors modestly claim that while their findings demonstrate the feasibility of using calibration models to generate realistic rollouts and recover meaningful hyperparameter trends, significant practical challenges remain. Specifically, the paper highlights open questions regarding:

The conditions under which extrapolating models fail.
How to effectively simulate distribution shifts associated with specific target deployment periods.
The necessity of predictive accuracy versus the preservation of hyperparameter rankings.
The development of metrics more informative than DTW for assessing long-horizon rollout quality.

The work bridges the gap between simulated calibration model research and real-world deployment, emphasizing that while the approach is promising, further investigation is required to address scalability and non-stationarity in complex industrial systems.