Source-linked AI summary
Online Reinforcement Learning in the Met Office Unified Model through Distributed Model-Agent Coupling
Pritthijit Nath, Sebastian Schemm, Peter Haynes, Emily Shuckburgh, Mark Webb
TL;DR
The paper asks whether machine-learnt corrections can adapt within an evolving global NWP model without sacrificing coupled stability. It couples the UKMO Unified Model to distributed DDPG agents, trains with analysis nudging, and evaluates a frozen policy without nudging. The workflow transfers successfully in the evaluated case, with selected Z500 and MSLP error reductions, but the single-case multivariable response is mixed.
Problem
Offline ML corrections trained on fixed datasets may not adapt to evolving model states and can destabilise coupled forecasts.
Method
The study couples the UKMO Unified Model to distributed RL agents that apply bounded potential-temperature corrections, using analysis nudging during training and frozen non-nudged inference.
Results
At +6 h, Z500 MAE reductions reached 45.8% and 40.8% in the northern and southern tropics, while the largest MSLP MAE reduction was 27.3% at 0–30°N.
Takeaways & Limitations
The successful transfer and short-horizon stability demonstrate promise for distributed online learning and RL-based bias correction within an operational NWP setting.
Takeaways & Limitations
This single-case experiment cannot establish transfer across regimes, and its multivariable response is mixed.
Abstract
from arXiv · showhide
Machine-learnt corrections can complement numerical weather prediction only if they adapt to the evolving model state while preserving dynamical consistency and numerical stability. To test this within a global forecasting model, we couple the Met Office (UKMO) Unified Model (UM) with distributed RL agents through rank-local tensors. A DDPG actor shares weights across the 70 vertical model levels of each atmospheric column and applies bounded potential-temperature corrections to the model tendencies. Across ten nudged training forecasts, nudging calculations towards the UKMO operational analysis provides an immediate counterfactual target. The frozen policy is then evaluated in a non-nudged forecast for inference. The coupled workflow successfully completes training and remains numerically stable in the evaluated case. Relative to a matched native UM forecast at +6 h, the learnt policy reduces Z$_{500}$ MAE in four of six latitude bands, including reductions of 45.8% and 40.8% in the northern and southern tropics. MSLP error too decreases in three bands, with a maximum reduction of 27.3% at 0-30°N. This single-case experiment demonstrates significant promise and feasibility of distributed online learning followed by non-nudged inference, laying the groundwork for RL-based bias correction and parametrisations within operational systems.
1 Introduction
The paper addresses the challenge of adapting machine-learnt corrections to evolving model states while preserving coupled numerical behaviour. It introduces distributed online RL in the UKMO Unified Model, including analysis-informed training and frozen non-nudged inference.
- Offline ML corrections may perform well on fixed data yet destabilise coupled numerical forecasts, motivating online adaptation.Online learning can adapt corrections to evolving states and under-represented regimes while remaining embedded in model constraints.
- RL maps the current model state to actions and optimises feedback from the resulting model state.The paper uses RL to pursue state-dependent corrections within an established global forecast model.
- The work demonstrates, to the authors’ knowledge, the first online RL execution within a distributed operational global NWP model.The coupling spans MPI/OpenMP Fortran model code and Python agents exchanging distributed fields.
- The system couples UM ranks and RL agents through rank-local tensors while respecting the model’s domain decomposition.The workflow also preserves an analysis-informed training target for later non-nudged inference.
- Training uses operational analysis nudging as an immediate counterfactual target, followed by deployment of a frozen policy without analysis nudging.Evaluation includes learning progression, numerical stability, and meteorological verification metrics.
2 Coupled-learning method
The coupled method distributes UM state, RL actions, and rewards across MPI-matched agents, using a shared DDPG actor to apply bounded potential-temperature corrections. Training rewards reductions in residual nudging requirements, while inference disables nudging and learning.
- Distributed coupling: The UM N320 atmosphere-only configuration uses 70 vertical levels, a 720-s exchange interval, and 192 MPI tiles matched by Python agents.Each agent represents the grid-cell columns belonging to its corresponding tile and rank.
- Distributed coupling: At each exchange, UM publishes a local state tensor, the agent returns an action tensor, and UM advances the forecast before publishing reward diagnostics.SmartRedis through SmartSim handles rank-local tensor exchange, while Cylc coordinates the workflow.
- Policy and action: Each column state concatenates five 70-level profiles into a 350-element vector, which the shared actor reshapes into level-wise five-feature inputs.The actor applies the same network at every active level and outputs actions in [-1, 1].
- Policy and action: The actor applies bounded potential-temperature increments to levels 2–69, with levels 1 and 70 inactive.The action is converted into a temperature-tendency correction using the stated scale factor.
- Reward design: The policy is evaluated by reducing the remaining nudging requirement rather than reproducing the analysis state directly.The native increment is the counterfactual target, and the RL correction changes the model state before the residual increment is diagnosed.
- Reward design: The reward combines a predominantly uniform vertical score with a smaller layer-mass-weighted contribution: 75% uniform and 25% mass-weighted.A positive column reward indicates weighted improvement across the column without requiring every level to improve.
- Train-to-inference transfer: During inference, nudging, exploration, replay insertion, and weight updates are disabled, leaving only the frozen actor to modify the forecast.The analysis remains available only for counterfactual diagnostics and evaluation.
3 Results
Ten nudged training forecasts produced a rising reward, and the episode-10 checkpoint transferred to a stable frozen non-nudged inference forecast. Against a native UM control at +6 h, Z500 and MSLP errors improved in selected latitude bands.
- Learning progression: 1.47×10^-3 was the mean reward in training forecast ten, up from 2.45×10^-4 in forecast one; inference remained close at 1.28×10^-3.The workflow completed ten nudged training forecasts and one frozen non-nudged inference forecast.
- Learning progression: Successful checkpoint restoration with nudging, exploration, and learning disabled demonstrated short-horizon stability during non-nudged inference.The forecasts were initialised at 00 UTC on 12 December 2021 and run to 6 hr 12 min.
- Forecast verification: 45.8% and 40.8% reductions in Z500 MAE occurred in the northern and southern tropics, respectively, relative to the native no-RL control.Z500 MAE decreased throughout 60°N–60°S.
- Forecast verification: 27.3% was the largest MSLP MAE reduction, occurring at 0–30°N relative to the native no-RL control.MSLP MAE decreased in three latitude bands.
4 Discussion and conclusion
The experiment demonstrates online RL execution and transfer to non-nudged inference, but its multivariable response is mixed and its single-case design cannot establish transfer across regimes.
- Polar Z500 MAE increased by 19.4% in the north and 18.9% in the south, while T1.5m and T500 MAE increased in every band.These opposing regional responses motivate regionally specialised agents.
- The single-case experiment cannot establish transfer across regimes.The conclusion is bounded by the tested initialisation and forecast case.
- Because local reward uses MSE while verification uses MAE, positive reward can coexist with broader absolute-error degradation.The potential-temperature objective may also leave other fields unconstrained.
- The authors identify seasons, seeds, matched controls, overhead measurements, and energy and mass budgets as future evaluation requirements.These additions are proposed to separate case-specific variability and assess broader operational behaviour.
Appendix A Implementation details
The implementation uses a DDPG actor with transformed five-profile, 70-level column states and bounded actions, distributed across a 12 × 16 MPI decomposition.
- The state contains potential temperature, one-step tendency, model level, pressure, and vertical-gradient profiles across 70 model levels.
- The 640 × 480 grid is distributed over 192 ranks in a 12 × 16 processor layout, with each rank owning 30 rows and 53 or 54 columns.
Appendix B Coupling workflow
The workflow coordinates forecast-level tasks with Cylc, timestep tensor exchange through Redis and SmartSim, and checkpoint persistence across training and inference.
- Cylc coordinates UM, Redis, and RL-agent tasks for each forecast, while checkpoints carry learned state between successive forecasts.
- SSDB mapping records the Redis host and port under the current task name, allowing training and inference tasks to resolve changing service locations.
- Within each timestep, matching ranks exchange local state, actions, diagnostic increments, and rewards through a synchronized state–action–diagnostic sequence.
- Training persists actor, critic, target networks, optimizer state, replay buffer, metadata, and global step; inference restores the final actor with updates disabled.
Appendix C Coupling profiling metrics
Profiling measures rank-local Python and selected-rank UM timings for the coupled workflow, showing substantial synchronization and waiting costs while limiting overhead comparisons.
- State retrieval and synchronization account for 49.9% of final-training agent wall time and 52.1% of inference agent wall time.
- Policy evaluation contributes 12.5% during training and 15.6% during inference, while 16 gradient updates contribute 4.7% during training.
- On UM rank 96, waiting for the inference action accounts for 28.9% of the 197-second UM task, while other UM activity accounts for 65.7%.
- The inclusive state-retrieval timer includes waiting for the next UM state, making it larger than the raw tensor GET time.
- The selected-rank trace characterizes the coupled execution boundary but does not quantify variation among UM ranks or matched overhead against native forecasts.
Appendix D Additional spatial verification
Appendix D compares native and coupled RL forecasts against UKMO analysis at +6 h using latitude-weighted errors and biases across four meteorological variables.
- Z500 MAE decreases in the four latitude bands between 60°N and 60°S, while T1.5m and T500 MAE increase in every band.
- The Z500 panels compare native control, coupled RL forecast, and UKMO analysis, then show absolute-error and bias differences relative to native control.
- T500 MAE increases in every latitude band, providing a contrasting diagnostic to the Z500 improvements.
- MSLP changes are spatially heterogeneous, consistent with latitude-band improvements and degradations across the verification domain.
- T1.5m responses are spatially mixed, while latitude-band MAE increases in all six regions.