arXiv CS open access

Online Reinforcement Learning in the Met Office Unified Model through Distributed Model-Agent Coupling

2026 · arxiv_cs

arXiv CS · Papers · 2026

deep-learning machine-learning neural-networks

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

Online Reinforcement Learning in the Met Office Unified Model through Distributed Model–Agent Coupling Pritthijit Nath1 Sebastian Schemm1 Peter Haynes1 Emily Shuckburgh2 Mark Webb3

arXiv:2609.02566v1 [cs.LG] 2 Sep 2026

Department of Applied Mathematics and Theoretical Physics, University of Cambridge 2 Department of Computer Science and Technology, University of Cambridge 3 Met Office Hadley Centre {pn341,ss3299,phh1,efs20}@cam.ac.uk; mark.webb@metoffice.gov.uk

Abstract 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 ranklocal 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 Z500 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.

Introduction

Artificial intelligence (AI) is increasingly used to accelerate weather prediction and represent processes that are difficult to resolve explicitly. Data-driven forecasting systems such as Aurora [1], ESFM [2], FourCastNet [3], Pangu-Weather [4], and GraphCast [5] produce skilful forecasts at substantially lower inference cost than conventional numerical weather prediction (NWP) methods. Hybrid systems instead embed machine-learning (ML) components within NWP, retaining the resolved dynamics while learning selected unresolved tendencies. NeuralGCM [6], for example, demonstrates that learnt components can operate stably within evolving weather and climate model dynamics, with similar stability demonstrated in multi-week ICON aquaplanet integrations under appropriate physical constraints [7]. Such approaches could make kilometre-scale simulations and larger ensembles more affordable while improving the representation of unresolved processes. However, their credibility ultimately depends on remaining stable and physically consistent as external forcing pushes the climate beyond the training distribution. Most ML-based corrections are learnt offline from fixed datasets and subsequently inserted into the numerical model without further adaptation. Strong offline performance therefore does not guarantee reliable coupled behaviour, and learnt corrections can destabilise the resolved state [7–9], although physical constraints and targeted treatments can improve coupled stability [7, 10]. Online learning therefore offers a promising alternative, allowing corrections to adapt continually to the evolving model state and to regimes under-represented in the original training data while remaining embedded within the physical constraints of the numerical model. Preprint. Under review.

Reinforcement Learning (RL) provides one potential route to online learning because a policy (a learnt mapping from the current model state to an action) can act on the evolving model and optimise feedback from the resulting state. Idealised weather and climate experiments [11] show that RL can learn state-dependent corrections and parametrisation controls. Extending this approach to an established global model creates a design challenge where MPI/OpenMP Fortran code must synchronise with Python agents, exchange distributed three-dimensional fields, and transfer a learnt policy from analysis-informed training to non-nudged operational inference where future analyses are unavailable. We address this challenge by coupling the Met Office (UKMO) Unified Model (UM) [12], an operational forecasting system, to Python-based RL agents through SmartSim [13], an in-memory interface between simulations and ML services at high-performance-computing scale. To the best of our knowledge, this work demonstrates the first execution of online RL within the distributed execution of an operational global NWP model. The key contributions of this work are: 1. An MPI rank-local online coupling that exchanges UM state, actions, and rewards with distributed RL agents through SmartSim while respecting the model’s domain decomposition. 2. A train-to-inference formulation in which operational analysis nudging provides an immediate counterfactual target during training while the frozen policy is subsequently deployed without analysis nudging during inference. 3. A coupled feasibility evaluation using learning progression, numerical stability, and standard meteorological verification metrics: latitude-weighted mean absolute error (MAE) in near-surface temperature, mean sea-level pressure, and 500-hPa temperature and geopotential height.

Coupled-learning method

2.1

Distributed model–agent exchange

We use an atmosphere-only UM N320 configuration with 40 km (approximate) horizontal resolution, 70 vertical model levels, and a 720-s timestep that also sets the model–agent exchange interval. The 640 × 480 horizontal grid is decomposed over a 12 × 16 = 192 tile layout, with each tile contained within a single MPI rank (shown in Figure A.1). Each tile contains 30 × 54 = 1, 620 atmospheric grid-cell columns. A matching set of Python agents mirrors this decomposition, with one agent assigned to each tile/rank and the corresponding grid-cell columns represented as vectorised environments. This 1-to-1 mapping avoids gathering distributed tiles onto a single process, since each UM rank exchanges only its local tile with the corresponding agent. At each exchange, a UM rank writes its local state tensor to Redis through SmartRedis, the client library within SmartSim. The corresponding Python rank evaluates the shared policy for every column and returns an action tensor with the same shape. The UM applies the action (corresponding to a temperature tendency), advances the forecast, and later publishes the reward diagnostics (defined in Section 2.3). Cylc [14] coordinates the entire workflow (Figure B.1) including the Redis server, the UM and agent tasks, ten training forecasts, checkpoint persistence, and the final inference forecast, with the associated model–agent exchange overheads quantified in Appendix C. 2.2

RL setup

For a grid-cell column, the state at level k contains normalised potential temperature θ̂k , its normalised c k , normalised model level k̂, vertical gradient g θ , and normalised pressure one-step tendency ∆θ k p̂k . Concatenating these five 70-level profiles gives a 350-element column state. The shared DDPG actor [15] reshapes this state into 70 five-feature vectors and applies the same network to every level, producing a fraction ak ∈ [−1, 1]. Levels 1 and 70 are inactive, while levels 2–69 receive the increment: δθkRL = θk fθ ak , fθ = 2.5 × 10−5 . (1) Figure 1 assesses the RL correction by comparing the nudging required before and after the policy action. The native increment nk = ∆θknative represents the correction from the uncorrected forecast towards the analysis-nudged state, while the residual increment ek = ∆θkresidual shows the correction still required after applying δθkRL . The policy is therefore evaluated by how much it reduces the remaining nudging requirement, rather than by reproducing the analysis state directly.

Figure 1: Potential-temperature increments at level k and exchange time t. Black nodes denote the model states model model RL θk,t and θk,t+1 , the RL-adjusted state θk,t+1 , and the analysis analysis θk,t+1 . The black and red arrows show the uncorrected model evolution and total model-to-analysis displacement respectively. The green native increment native ∆θk,t provides the counterfactual target. The magenta RL RL correction δθk,t changes the model state before nudgresidual ing, leaving the blue residual increment ∆θk,t . Arrow lengths and directions are illustrative.

Each MPI-matched agent rank holds its own actor, critic, target networks, and replay buffer. The actor and critic share their respective weights across all active levels and columns on that rank, and each has two 512-unit hidden layers. Level-wise transitions are assigned to five pressure strata separated at 200, 400, 600, and 850 hPa, from which each batch draws approximately equal numbers of samples. Because the target HPC partition is CPU-only, the agent tasks run on CPUs. The replay capacity is 500,000 transitions, the batch size is 256, and each eligible exchange performs 16 gradient updates. Table A.1 gives the complete numerical configuration, while Algorithm 1 describes the update sequence. 2.3

Analysis-informed reward and non-nudged inference

For active levels A = {2, . . . , 69}, we define the layer mass mk = ∆pk /g, where ∆pk is the layerpressure thickness P and g is gravitational acceleration. The mean-normalised value is m e k = mk /mA with mA = |A|−1 k∈A mk being the mean layer mass. The level score and column reward are: sk =

n2k − e2k , 2 nk + e2k + ϵ32

1 X wk sk , |A|

wk = cu + cm m e k,

cu = 0.75,

cm = 0.25.

k∈A

(2) The constant ϵ32 prevents division by zero in single-precision arithmetic. As |n2k − e2k | ≤ n2k + e2k (nk and ek defined in Section 2.2) and the denominator is larger by ϵ32 , |sk | < 1. A positive score means that the RL action reduces the squared residual correction at that level. For no-RL situations, nk = ek resulting in sk = 0. The coefficients cu and cm combine a predominantly uniform vertical mean (75%) with a smaller layer-mass-weighted contribution (25%). A positive column reward therefore indicates a weighted improvement across the column without requiring every level to improve. During training, the UM applies the RL correction before nudging, as shown in Figure B.1. By contrast during inference, nudging is disabled, although the same counterfactual target is diagnosed for evaluation. The analysis is neither supplied to the actor nor applied to the prognostic state. Exploration, replay insertion, and weight updates are likewise disabled, leaving only the frozen actor to modify the forecast.

Results

3.1

Transfer to non-nudged execution

The coupled workflow completed ten nudged training forecasts and one frozen non-nudged inference forecast, all initialised at 00 UTC on 12 December 2021 and run to a lead time of 6 hr 12 min. Each forecast contained 31 model–agent exchanges at 12-min intervals. Aggregating one value from each of the 192 ranks at every exchange therefore generated 192 × 31 = 5,952 rank-time reward summaries per forecast. As Figure 2 shows, the mean reward, averaged over all agents and exchange times, increased overall from 2.45 × 10−4 in the first training forecast to 1.47 × 10−3 in the tenth, and remained close to the final training value at 1.28 × 10−3 during inference. Restoring the distributed checkpoint with nudging, exploration, and learning disabled demonstrated successful transfer and short-horizon stability under interactive feedback.

Mean reward (×10−4)

12 10 Nudged training Non-nudged inference No RL

8 6 4

Figure 2: Learning progression curve. The ordinate is the rank-time mean reward r̄ expressed as r̄/10−4 . Red circles denote the ten connected nudged training episodes, the blue triangle denotes frozen non-nudged inference from the episode-10 checkpoint, and the dashed line denotes the no-RL reward of zero.

2 0 1

Training episode

Table 1: Latitude-band forecast-error changes relative to the native no-RL control. Each cell reports latitude-weighted MAE at +6 h against UKMO analysis, with absolute and percentage changes relative to native MAE. Negative values indicate lower error in the coupled RL forecast. Latitude bands are area-weighted grid-cell aggregates. Near-surface diagnostics Latitude 60–90◦ N 30–60◦ N 0–30◦ N 0–30◦ S 30–60◦ S 60–90◦ S

3.2

◦

T1.5m ( C)

500-hPa diagnostics ◦

T500 ( C)

MSLP (hPa)

Z500 (m)

∆MAE

Change (%)

∆MAE

Change (%)

∆MAE

Change (%)

∆MAE

Change (%)

+0.002 +0.051 +0.029 +0.024 +0.024 +0.013

+0.1% +7.4% +6.9% +7.1% +10.3% +3.0%

+0.051 −0.015 −0.129 −0.035 +0.012 +0.043

+12.7% −2.7% −27.3% −7.8% +3.5% +6.6%

+0.013 +0.005 +0.025 +0.038 +0.002 +0.015

+2.2% +0.8% +3.4% +4.9% +0.4% +2.5%

+0.402 −0.335 −1.822 −1.578 −0.266 +0.772

+19.4% −10.9% −45.8% −40.8% −8.5% +18.9%

Forecast comparison across variables and latitude

As the scalar reward does not necessarily reflect forecast quality across all variables, we also evaluate 1.5-m air temperature (T1.5m ), mean sea-level pressure (MSLP), 500-hPa temperature (T500 ), and 500-hPa geopotential height (Z500 ). We compare the +6 h coupled RL forecast with a native nonnudged UM control from the same initial state, using UKMO analysis as the reference. Table 1 reports latitude-weighted changes in MAE, while Appendix D shows the corresponding spatial fields, absolute errors, and biases. The Z500 MAE decreases throughout 60◦ N–60◦ S, with reductions of 1.822 m (45.8%) and 1.578 m (40.8%) in the northern and southern tropics. The MSLP MAE decreases in three bands, with the largest reduction of 0.129 hPa (27.3%) at 0–30◦ N.

Discussion and conclusion

Building on Nath et al. [11], this study demonstrates the first execution of online RL within a global NWP model. Fortran UM ranks and RL agents exchanged decomposed tensors, preserved learning state across forecasts, and restored the final checkpoint for frozen, non-nudged inference. The diagnosed analysis increment retains the same counterfactual meaning in both phases, with inference using the analysis only for diagnostics and verification. The multivariable response is mixed: polar Z500 MAE increases by 0.402 m (19.4%) in the north and 0.772 m (18.9%) in the south, while T1.5m and T500 MAE increase in every band. These opposing regional responses motivate regionally specialised agents, although this single case cannot establish transfer across regimes. The local reward uses MSE while verification uses MAE, so reducing a few large errors can improve reward despite broader absolute-error degradation. A positive local reward therefore does not imply broader forecast improvement, and the potential-temperature objective may leave other fields unconstrained. With one initialisation and no seed ensemble, these effects cannot be separated from case-specific variability. Future work would span seasons and seeds with matched native and nudging controls, overhead measurements, and energy and mass budgets. Within the current evidence boundary, the successful transfer of a state-dependent online policy to non-nudged inference in a distributed NWP setting demonstrates significant promise and provides a practical basis for evaluating RL-derived bias correction and parametrisations of unresolved sub-grid processes.

Acknowledgements and Disclosure of Funding P. Nath was supported by the UKRI Centre for Doctoral Training in Application of Artificial Intelligence to the study of Environmental Risks [EP/S022961/1]. Mark Webb was supported by the Met Office Hadley Centre Climate Programme funded by DSIT. The authors thank the Met Office for providing access to the Unified Model, analysis data, and the Met Office EX computing platform under a CASE studentship agreement with the University of Cambridge. We thank Andrew Shao and Alessandro Rigazzi (HPE) for their valuable assistance with setting up SmartSim.

References [1] Bodnar C, Bruinsma WP, Lucic A, et al. A foundation model for the Earth system. Nature. 2025;641:1180-7. Available from: https://doi.org/10.1038/s41586-025-09005-y. [2] Ozdemir F, Cheng Y, Mohebi S, Lehmann F, Adamov S, Zhang Z, et al. Earth System Foundation Model (ESFM): A unified framework for heterogeneous data integration and forecasting. arXiv preprint arXiv:260500850. 2026. Available from: https://arxiv.org/abs/2605.00850. [3] Pathak J, Subramanian S, Harrington P, et al. FourCastNet: A global data-driven high-resolution weather model using adaptive Fourier neural operators. arXiv preprint arXiv:220211214. 2022. Available from: https://arxiv.org/abs/2202.11214. [4] Bi K, Xie L, Zhang H, Chen X, Gu X, Tian Q. Accurate medium-range global weather forecasting with 3D neural networks. Nature. 2023;619:533-8. Available from: https://doi.org/10.1038/ s41586-023-06185-3. [5] Lam R, Sanchez-Gonzalez A, Willson M, et al. Learning skillful medium-range global weather forecasting. Science. 2023;382:1416-21. Available from: https://doi.org/10.1126/science.adi2336. [6] Kochkov D, Yuval J, Langmore I, et al. Neural general circulation models for weather and climate. Nature. 2024;632:1060-6. Available from: https://doi.org/10.1038/s41586-024-07744-y. [7] Bertoli G, Mohebi S, Ozdemir F, Jucker J, Rüdisühli S, Perez-Cruz F, et al. Revisiting machine learning approaches for short- and longwave radiation inference in weather and climate models. Journal of Advances in Modeling Earth Systems. 2025;17(9):e2025MS004956. Available from: https://doi.org/10.1029/ 2025MS004956. [8] Rasp S, Pritchard MS, Gentine P. Deep learning to represent subgrid processes in climate models. Proceedings of the National Academy of Sciences. 2018;115:9684-9. Available from: https://doi. org/10.1073/pnas.1810286115. [9] Brenowitz ND, Bretherton CS. Spatially extended tests of a neural network parametrization trained by coarse-graining. Journal of Advances in Modeling Earth Systems. 2019;11:2728-44. Available from: https://doi.org/10.1029/2019MS001711. [10] Yuval J, O’Gorman PA. Use of neural networks for stable, accurate and physically consistent parameterization of subgrid atmospheric processes with good performance at reduced precision. Geophysical Research Letters. 2021;48:e2020GL091363. Available from: https://doi.org/10.1029/2020GL091363. [11] Nath P, Schemm S, Moss H, Haynes P, Shuckburgh E, Webb MJ. Replacing tunable parameters in weather and climate models with state-dependent functions using reinforcement learning. Journal of Advances in Modeling Earth Systems. 2026;18(8):e2026MS005745. Available from: https://doi.org/10.1029/ 2026MS005745. [12] Brown A, Milton S, Cullen M, Golding B, Mitchell J, Shelly A. Unified modeling and prediction of weather and climate: A 25-year journey. Bulletin of the American Meteorological Society. 2012;93:1865-77. Available from: https://doi.org/10.1175/BAMS-D-12-00018.1. [13] Partee S, Ellis M, Rigazzi A, et al. Using machine learning at scale in HPC simulations with SmartSim: An application to ocean climate modeling. Journal of Computational Science. 2022;62:101707. Available from: https://doi.org/10.1016/j.jocs.2022.101707. [14] Oliver H, Shin M, Matthews D, Sanders O, Bartholomew S, Clark A, et al. Workflow automation for cycling systems. Computing in Science & Engineering. 2019;21(4):7-21. Available from: https: //doi.org/10.1109/MCSE.2019.2906593.

[15] Lillicrap TP, Hunt JJ, Pritzel A, et al. Continuous control with deep reinforcement learning. International Conference on Learning Representations. 2016. Available from: https://arxiv.org/abs/1509. 02971. [16] Saarinen S, Hamrud M, Salmond D, Hague J. Dr. Hook instrumentation tool. Tech. Rep., European Centre for Medium-Range Weather Forecasts; 2005. Available from: https://www.umr-cnrm.fr/surfex/ IMG/pdf/DrHook.pdf.

Appendix A

Implementation details

Algorithm 1 summarises the DDPG updates used to train the actor and critic shared across the active levels and columns on each agent rank. Table A.1 collects the associated state transformations, action range, and learning settings. The state comprises potential temperature θk , one-step tendency ∆θk , model level k, pressure pk , and vertical gradient gkθ , from which the actor returns action ak . Potential temperature is mapped logarithmically, while tendency, model level, pressure, and vertical gradient use linear transformations. Pressure is bounded by pref = 1,100 hPa, which lies above standard sea-level pressure and provides a fixed upper bound for the model-level pressure. The same reference defines the dimensionless upward coordinate − log(p/pref ) used to calculate the vertical gradient from normalised potential temperature. This coordinate is zero at pref and increases as pressure decreases, while a 0.01-hPa floor keeps the logarithm finite at the model top. Table A.1: State transformations, action range, and DDPG settings. All five state profiles contain 70 model levels. Quantity

Physical or numerical range

Agent representation

Potential temperature θk One-step tendency ∆θk Model level k Pressure pk Vertical gradient gkθ Action ak Actor and critic Replay and updates

150–6000 K −10–10 K 1–70 0–1,100 hPa clipped to [−1, 1] [−1, 1] two hidden layers each capacity 500,000

logarithmically mapped to [−1, 1] linearly mapped to [−1, 1] linearly mapped to [−1, 1] linearly mapped to [−1, 1] unchanged after clipping levels 2–69 active 512 units per layer batch 256, 16 updates

The coupling follows the UM domain decomposition shown in Figure A.1. Each agent rank operates on the atmospheric columns owned by one UM rank, so policy evaluation and reward construction remain local and no process assembles a global field. This arrangement preserves the model’s distributed-memory layout while exposing every column in a tile as a member of the agent’s synchronous vector environment. Algorithm 1 Level-wise DDPG training with a shared actor and critic 1: Input: active levels A, pressure strata P , discount γ, target coefficient τ , batch size B, update

count G, and exploration noise ϵ 2: Initialise: shared actor πθ : R5 → [−1, 1], shared critic Qϕ : R5 × R → R, target networks πθ′

and Qϕ′ , and pressure-stratified replay {Dj }j∈P

3: for each model–agent exchange t do 4: Receive column states xc ∈ R5×70 for all local columns c 5: for all c and k ∈ A do 6: Select ac,k = πθ (xc,k ) + ϵc,k and apply all level actions together 7: Observe x′c,k , level reward rc,k , and termination flag dc 8: Store (xc,k , ac,k , rc,k , x′c,k , dc ) in the stratum Dj selected by pc,k 9: end for 10: if t > learning_starts then 11: for g = 1 to G do 12: Sample B transitions approximately equally across populated strata ′ 13: yi ← ri + γ(1 − di )Qϕ′ (x′i , πP θ ′ (xi )) −1 14: Update ϕ by minimising B Pi (Qϕ (xi , ai ) − yi )2 15: Update θ by maximising B −1 i Qϕ (xi , πθ (xi )) at the policy-update frequency 16: Soft-update each target ψ ′ ← τ ψ + (1 − τ )ψ ′ for ψ ∈ {θ, ϕ} 17: end for 18: end if 19: end for

UM Atmospheric MPI Decomposition | 12 × 16 (192 Ranks)

60°N

30°N

0° 8

30°S

60°S

180°

180

181

182

183

184

185

186

187

188

189

190

191

168

169

170

171

172

173

174

175

176

177

178

179

156

157

158

159

160

161

162

163

164

165

166

167

144

145

146

147

148

149

150

151

152

153

154

155

132

133

134

135

136

137

138

139

140

141

142

143

120

121

122

123

124

125

126

127

128

129

130

131

108

109

110

111

112

113

114

115

116

117

118

119

100

101

102

103

104

105

106

107

120°W

60°W

0°

60°E

120°E

180°

X: West-to-East Processor Column | Y: South-to-North Processor Row | Rank = (Y × 12) + X Figure A.1: Horizontal MPI decomposition of the 640 × 480 N320 atmosphere grid. The 12 × 16 processor layout contains 192 ranks, numbered from south to north and west to east. Each rank owns 30 rows and either 53 or 54 columns, giving 1,590 or 1,620 atmospheric columns. A single Python agent rank manages the vectorised environments for the each tile.

Appendix B

Coupling workflow

Figure B.1 shows how the workflow separates forecast orchestration, timestep communication, and learning-state persistence. Cylc controls forecast-level task dependencies, Redis handles timesteplevel tensor exchange between the UM and agents, and checkpoints carry the learnt state between successive forecasts. Training Episode K

RL Agent Stochastic Exploration w/ Weight Updates

Retrieve UKMO Analysis

UM Reconfigure, Map and Prepare initial conditions

Store checkpoint at regular intervals

UM Forecast with nudging

Perform Postprocesing

Archive into long-term storage (e.g., MASS)

Compile UM + Model Control Drivers using FCM Make Redis Server

Checkpoint Store

Load final checkpoint

Inference Episode RL Agent Deterministic w/o Weight Updates

Housekeeping

Archive into long-term storage (e.g., MASS)

Perform Postprocesing

UM Forecast without nudging

Redis Server

Figure B.1: Cylc workflow for training and inference. Each nudged training forecast runs the UM, Redis server, and RL-agent tasks together and writes the learning state to the checkpoint store. The inference branch restores the final checkpoint, disables exploration and weight updates, and runs the UM without applying nudging. Post-processing and archival remain outside the model–agent exchange. The PBS scheduler can place the Redis service on a different host for each forecast. To avoid embedding a database address in either coupled component, the workflow records the assigned host and port in the shared SSDB mapping. The UM and agent then resolve the address independently using the current task name. The task name changes between training and inference, but service discovery and the tensor protocol do not (Figure B.2). Communication is synchronised at two timescales. Within each forecast, exchanges follow a state– action–diagnostic sequence: every UM rank blocks until its matching agent consumes the rank-local state and returns an action, after which the UM publishes the native and residual increments used to evaluate that action. Rank-qualified keys identify each field and its owning rank, while delete-afterconsumption polling prevents actions or diagnostics from being reused at later timesteps. Together, these mechanisms form a barrier that prevents the UM from advancing with a stale action and prevents the agent from associating diagnostics with the wrong state. Figure B.4 illustrates the general protocol for tensor exchange; the theta-nudging experiment uses different variables while preserving the same ordering and rank-local ownership. Across forecasts, the workflow persists the actor, critic, target networks, optimiser, replay buffer, metadata, and global step, as shown in Figure B.3. Subsequent training forecasts therefore resume the 9

accumulated optimisation trajectory rather than restarting, whereas inference restores the final actor with exploration, replay insertion, and gradient updates disabled, making each action a deterministic function of the saved policy and current UM state. 2. Write to the SSDB mapping file on Redis database initialisation

4. Communication via SmartSim

1. Initialise the Redis server

RL Agent

TRAINING UM_FORECAST INFERENCE UM_FORECAST_INF

UM task CYLC_TASK_NAME = um_forecast

SSDB=HOST:PORT

HOST:PORT HOST:PORT

3. Extract SSDB mapping and set as REDIS_ADDRESS or ENVIRONMENT VARIABLE Shared SSDB mapping file

Figure B.2: Redis discovery through the shared SSDB mapping. The workflow starts the Redis server and writes its host and port to the mapping file under the current Cylc task name. The UM and RL-agent tasks read this mapping, set the Redis address, and then communicate through SmartSim. Separate task names distinguish training from inference without changing the tensor protocol. The SSDB mapping prevents the UM and agents from depending on a fixed database address. This indirection is required because the scheduler can place the Redis service on a different host for each workflow execution.

Checkpoint Writer RL Agent TRAINING with weight updates

Periodic Save (at regular intervals)

Final Save

SIGTERM_REDIS

DDPG

UM TRAINING with nudging

Checkpoint Store

UM INFERENCE without nudging

actor, target actor weights critic, target critic weights optimizer states replay buffer metadata, global step

RL Agent INFERENCE without noise + updates

Figure B.3: Checkpoint transfer between training and inference. During training, the writer saves the actor, critic, target networks, optimiser state, replay buffer, metadata, and global step. The final training checkpoint is restored for inference, where exploration and parameter updates are disabled. The complete checkpoint allows learning to continue across training forecasts and makes inference a deterministic evaluation of the frozen policy rather than a newly initialised agent.

Met Office Unified Model

FedRAIN - UM GridOrchestrator + DDPG

Redis Database

1a. Send the state information to the Redis database 1b. Retrieve the state information from the Redis database

PUT state.rank.<r>

GET state.rank.<r>

2b. Send the policy-driven prediction to the Redis database

PUT action.rank.<r>

2c. Retrieve the policy-driven prediction from the Redis database GET action.rank.<r> 3a. Insert the policy-driven prediction into the UM and perform the next model integration step + Calculate the nudge against a reference (UKMO analysis)

2a. Generate a policy-driven prediction as function of the state

3b. Send the temperature nudge to the Redis database

3c. Retrieve the temperature nudge from the Redis database

PUT nudge.rank.<r>

GET nudge.rank.<r>

4. Calculate the reward (and optionally, update the RL-agent under training)

Figure B.4: Rank-local UM–agent exchange and reward sequence for MPI rank ⟨r⟩. (1) The UM publishes its local state which the matching DDPG agent retrieves. (2) The agent evaluates the policy and returns the predicted correction. (3) The UM retrieves and applies this action, advances the model by one integration step, diagnoses the temperature-nudging increment relative to the UKMO analysis, and publishes it. (4) The agent retrieves the increment to calculate the reward and, during training, optionally updates the policy. Redis mediates each PUT–GET transfer; rank-qualified keys preserve the UM domain decomposition and keep state, action, and reward diagnostics associated with the same model step.

Appendix C

Coupling profiling metrics

For profiling purposes, the experiment was instrumented at the Python model–agent boundary and within the UM. Table C.1 reports the mean and population standard deviation of those rank-local means across all 192 ranks. The final training and inference agent tasks used 4x Genoa nodes and took 242 and 198 seconds, respectively. The corresponding UM tasks used 192 MPI ranks with 2x OpenMP threads per rank on 2x Genoa nodes and took 237 and 197 seconds. These task durations include initialisation and shutdown but exclude scheduler queueing, post-processing, and archival. Table C.1: Rank-local Python wall-clock timings for the final nudged training forecast and frozen non-nudged inference. Values are the mean ± population standard deviation across 192 rank means, in seconds per invocation.1 Component

Final training Mean ± spread (s)

Inference Mean ± spread (s)

Full agent exchange State retrieval including synchronisation State tensor GET Policy evaluation Action tensor PUT Environment and reward step 16 gradient updates

5.887 ± 0.031 3.898 ± 0.082 0.0258 ± 0.0040 0.972 ± 0.014 0.0225 ± 0.0097 0.0617 ± 0.0010 0.377 ± 0.070

4.792 ± 0.070 3.326 ± 0.079 0.0241 ± 0.0033 0.999 ± 0.013 0.0263 ± 0.0119 0.0618 ± 0.0009 —

Figure C.5 partitions the complete agent-task wall time into mutually exclusive components. State retrieval and synchronisation account for 49.9% of final-training wall time and 52.1% of inference wall time. Policy evaluation contributes 12.5% and 15.6%, respectively, while the 16 gradient updates contribute 4.7% during training. The remaining agent categories are each below 8%, apart from the 18.1% residual outside the profiled exchange in both tasks. On a specific UM rank, waiting for the inference action accounts for 28.9% of the 197-s UM task, while state publication, action retrieval, and reward publication contribute 2.9%, 0.3%, and 2.2%. Other UM activity outside these four instrumented SmartRedis regions accounts for the remaining 65.7%. Final nudged training

Frozen non-nudged inference

12.5%

UM inference (single rank)

15.6% 28.9%

49.9%

242 s

52.1%

198 s

18.1%

State retrieval and synchronisation Policy evaluation Gradient updates Environment and reward

18.1%

Action tensor PUT Other exchange activity Initialisation Outside profiled exchange

197 s 65.7%

State publication Action wait Action retrieval Reward publication Other UM activity

Figure C.5: Share of complete task wall time for the final nudged training agent task, frozen nonnudged inference agent task, and inference UM task on rank PE96. Agent component durations are rank-mean cumulative timings, calculated as the mean time per invocation multiplied by the mean invocation count. Nested agent timers are not counted twice. The PE96 panel multiplies the Dr Hook mean by 31 state, wait, and action-retrieval calls and by 30 reward-publication calls. Other UM activity is the residual of the 197-s UM task. 1 The full-exchange and state-retrieval timers contain some of the other rows and therefore must not be added together.

The inclusive state-retrieval timer is much larger than the raw tensor GET because it includes waiting for the next UM state. Similarly, the full agent-exchange timer covers synchronisation, policy evaluation, tensor operations, reward construction, and training when enabled. The inference exchange is 1.095 s shorter on average. Inference omits gradient updates, while much of the remaining difference occurs in the inclusive state-retrieval wait and the policy-evaluation cost remains similar. The small standard deviations for the full exchange and policy evaluation show that these costs are balanced across the horizontal decomposition in this case. The profiling tool Dr. Hook [16] separately measures the Fortran side on UM rank 96. Table C.2 reports the mean time per invocation for the four instrumented SmartRedis regions during final training and inference. Table C.2: Dr. Hook wall-clock timings on UM rank 96 for the final nudged training forecast and frozen non-nudged inference. Component values are mean seconds per invocation; the final row gives the instrumentation overhead reported by Dr. Hook. Component

Final training Mean time (s)

Inference Mean time (s)

0.193 2.652 0.0169 0.138 1.12

0.185 1.835 0.0199 0.145 1.34

State publication Wait for agent action Action retrieval Reward-diagnostic publication Dr. Hook overhead (%)

This selected-rank trace identifies where the synchronous UM process waits, but it does not quantify variation among UM ranks. Consequently, the measurements characterise the coupled execution boundary rather than providing a matched estimate of overhead relative to an uninstrumented native forecast.

Appendix D

Additional spatial verification

T1.5m | MAE

60-90N

30-60N

0-30N

30S-0

60-30S

90-60S

90-60S 0.0

0.5

1.0

1.5

T1.5m | Bias

−1.0

−0.5

MAE (°C)

MSLP | MAE

60-90N

30-60N

0-30N

30S-0

60-30S

90-60S

90-60S 0.0

0.2

0.4

0.6

−0.2

−0.1

60-90N

30-60N

0-30N

30S-0

60-30S

90-60S

90-60S 0.2

0.4

0.6

0.8

−0.5

Z500 | MAE

60-90N

30-60N

0-30N

30S-0

60-30S

90-60S

90-60S 1

0.1

0.2

0.0

0.5

Bias (°C)

60-90N

0.0

T500 | Bias

MAE (°C)

1.0

Bias (hPa)

T500 | MAE

0.0

0.5

MSLP | Bias

MAE (hPa)

0.0

Bias (°C)

Z500 | Bias

−2

MAE (m)

Bias (m) Native

Figure D.1: Latitude-weighted zonal verification at +6 h. Panels a, c, e, and g compare native and RL mean absolute error (MAE) for T1.5m , MSLP, T500 , and Z500 , respectively. Panels b, d, f, and h show the corresponding signed biases. Each bar is an area-weighted grid-cell aggregate within one latitude band. The coupled response is mixed across variables and regions: Z500 MAE decreases in the four latitude bands between 60◦ N and 60◦ S, whereas T1.5m and T500 MAE increase in every band. These zonal summaries describe one forecast initialisation and are not independent samples.

Z500 | Lead time: +06h Native control

5000

5200

−10

5400

5600

5800

5000

5200

5400

5600

UKMO analysis

5800

5000

5200

5400

5600

MAE Native control

MAE Coupled RL forecast

MAE difference Coupled RL forecast - Native control

−5

Coupled RL forecast

−6

−4

−2

Bias Native control

Bias Coupled RL forecast

Bias difference Coupled RL forecast - Native control

−10

−5

−6

−4

−2

5800

Figure D.2: Z500 verification at +6 h. (a) The first row compares the native control, coupled RL forecast, and UKMO analysis. (b) The second row shows absolute error for each forecast and the RLminus-native absolute-error difference. (c) The third row shows forecast bias and the corresponding RL-minus-native bias difference. Blue values in the right-hand difference panels indicate a reduction relative to the native control, while red values indicate an increase. The Z500 MAE decreases through the four latitude bands between 60◦ N and 60◦ S, while both polar bands deteriorate. The spatial pattern provides context for the latitude-weighted summary, but it represents only one forecast initialisation.

T500 | Lead time: +06h Native control

−40

0.0

−30

−20

−10

−40

−30

UKMO analysis

−20

−10

−40

−30

−20

−10

°C

MAE Native control

MAE Coupled RL forecast

MAE difference Coupled RL forecast - Native control

0.5

−1

Coupled RL forecast

1.0

1.5

0.0

0.5

1.0

1.5

−0.6

−0.4

−0.2

0.0

0.2

°C

Bias Native control

Bias Coupled RL forecast

Bias difference Coupled RL forecast - Native control

°C

−1

°C

−0.6

−0.4

−0.2

0.0

0.2

0.4

0.6

0.4

0.6

°C

Figure D.3: T500 verification at +6 h. (a) The first row compares the native control, coupled RL forecast, and UKMO analysis. (b) The second row shows absolute error and its RL-minus-native difference. (c) The third row shows bias and its RL-minus-native difference. Blue values in the right-hand difference panels indicate a reduction relative to the native control, while red values indicate an increase. The T500 MAE increases in every latitude band, so this diagnostic provides an important counterexample to the Z500 improvement. Better geopotential-height verification therefore does not necessarily imply a corresponding improvement in temperature at the same pressure level.

MSLP | Lead time: +06h Native control

970

0.0

980

0.5

−2

990

1010

1020

1030

970

980

990

1000

1010

1020

UKMO analysis

1030

970

980

990

1000

1010

1020

hPa

MAE Native control

MAE Coupled RL forecast

MAE difference Coupled RL forecast - Native control

1.0

−1

1000

Coupled RL forecast

1.5

2.0

2.5

0.0

0.5

1.0

1.5

2.0

2.5

−0.75

−0.50

−0.25

0.00

0.25

hPa

Bias Native control

Bias Coupled RL forecast

Bias difference Coupled RL forecast - Native control

hPa

−2

−1

hPa

−0.5

0.0

0.50

0.5

hPa

Figure D.4: MSLP verification at +6 h. (a) The first row compares the native control, coupled RL forecast, and UKMO analysis. (b) The second row shows absolute error and its RL-minus-native difference. (c) The third row shows bias and its RL-minus-native difference. Blue values in the right-hand difference panels indicate a reduction relative to the native control, while red values indicate an increase. MSLP changes are spatially heterogeneous, consistent with the latitude-band improvements in three of six regions and degradation in the remaining three. The map describes this single initialisation and does not provide independent spatial samples.

1030

0.75

T1.5m | Lead time: +06h Native control

−20

−4

Coupled RL forecast

−20

°C

MAE Native control

MAE Coupled RL forecast

MAE difference Coupled RL forecast - Native control

−2

UKMO analysis

−1.0

−0.5

0.0

0.5

°C

Bias Native control

Bias Coupled RL forecast

Bias difference Coupled RL forecast - Native control

°C

−4

−2

°C

−1.0

−0.5

0.0

0.5

°C

Figure D.5: T1.5m verification at +6 h. (a) The first row compares the native control, coupled RL forecast, and UKMO analysis. (b) The second row shows absolute error and its RL-minus-native difference. (c) The third row shows bias and its RL-minus-native difference. Blue values in the right-hand difference panels indicate a reduction relative to the native control, while red values indicate an increase. The T1.5m response is also spatially mixed, yet latitude-band MAE increases in all six regions. Despite the gains in Z500 accuracy, near-surface degradation highlights a persistent misalignment between immediate rewards and overall multivariable forecast skill.

1.0

Online Reinforcement Learning in the Met Office Unified Model through Distributed Model-Agent Coupling

Analyze & explore related research in Conceptio

Related documents