feat(rl): add replay-buffer persistence for interrupted training runs

assume/common/base.py

@@ -867,6 +867,11 @@ class LearningConfig:
     tau: float = 0.005
     target_policy_noise: float = 0.2
     target_noise_clip: float = 0.5
+    # Replay buffer persistence
+    save_replay_buffer: bool = True
+    replay_buffer_save_path: str | None = None
+    load_replay_buffer: bool = False
+    replay_buffer_load_path: str | None = None
 
     def __post_init__(self):
         """Calculate defaults that depend on other fields and validate inputs."""

assume/reinforcement_learning/buffer.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: AGPL-3.0-or-later
 
+import os
 import warnings
 from typing import NamedTuple
 
@@ -174,3 +175,42 @@ def sample(self, batch_size: int) -> ReplayBufferSamples:
         )
 
         return ReplayBufferSamples(*tuple(map(self.to_torch, data)))
+
+    def save(self, path: str):
+        """Save the replay buffer state to disk."""
+        # ensure directory exists
+        dirpath = os.path.dirname(path)
+        if dirpath and not os.path.exists(dirpath):
+            os.makedirs(dirpath, exist_ok=True)
+
+        np.savez_compressed(
+            path,
+            observations=self.observations,
+            actions=self.actions,
+            rewards=self.rewards,
+            pos=np.array([self.pos]),
+            full=np.array([self.full]),
+        )
+
+    @classmethod
+    def load(cls, path: str, device: str, float_type):
+        """Load a replay buffer from disk."""
+        data = np.load(path)
+        obs = data["observations"]
+        acts = data["actions"]
+        rews = data["rewards"]
+
+        buffer = cls(
+            buffer_size=obs.shape[0],
+            obs_dim=obs.shape[2],
+            act_dim=acts.shape[2],
+            n_rl_units=obs.shape[1],
+            device=device,
+            float_type=float_type,
+        )
+        buffer.observations = obs
+        buffer.actions = acts
+        buffer.rewards = rews
+        buffer.pos = int(data["pos"][0])
+        buffer.full = bool(data["full"][0])
+        return buffer

assume/scenario/loader_csv.py

@@ -1167,7 +1167,8 @@ def run_learning(
     from assume.reinforcement_learning.buffer import ReplayBuffer
 
     if not verbose:
-        logger.setLevel(logging.WARNING)
+        # Avoid silencing this module's INFO logs; instead silence very noisy external loggers
+        logging.getLogger("mango").setLevel(logging.WARNING)
 
     # remove csv path so that nothing is written while learning
     temp_csv_path = world.export_csv_path
@@ -1188,15 +1189,63 @@ def run_learning(
 
     # -----------------------------------------
     # Information that needs to be stored across episodes, aka one simulation run
-    inter_episodic_data = {
-        "buffer": ReplayBuffer(
+    # Read optional replay-buffer persistence settings from learning_config
+    lc = world.learning_role.learning_config
+    cfg_save_flag = getattr(lc, "save_replay_buffer", True)
+    cfg_save_path = getattr(lc, "replay_buffer_save_path", None)
+    cfg_load_flag = getattr(lc, "load_replay_buffer", False)
+    cfg_load_path = getattr(lc, "replay_buffer_load_path", None)
+
+    # default path next to saved policies
+    default_buffer_path = f"{save_path}/last_policies/replay_buffer.npz"
+
+    buffer = None
+    # Load only when explicitly requested via load_replay_buffer
+    if cfg_load_flag:
+        # choose explicit load path if provided, otherwise fall back to configured save path or default
+        path_to_load = (
+            cfg_load_path
+            if cfg_load_path is not None
+            else (cfg_save_path if cfg_save_path is not None else default_buffer_path)
+        )
+        if not os.path.exists(path_to_load):
+            raise AssumeException(
+                f"load_replay_buffer is true but no buffer file found at {path_to_load}"
+            )
+        try:
+            buffer = ReplayBuffer.load(
+                path_to_load,
+                device=world.learning_role.device,
+                float_type=world.learning_role.float_type,
+            )
+            logger.info(f"Loaded replay buffer from {path_to_load}")
+            # disable initial experience collection when buffer provided
+            try:
+                world.learning_role.learning_config.episodes_collecting_initial_experience = 0
+                logger.info(
+                    "Replay buffer provided — skipping initial experience collection (episodes_collecting_initial_experience set to 0)."
+                )
+            except Exception:
+                logger.warning(
+                    "Could not set episodes_collecting_initial_experience to 0 on learning_config"
+                )
+        except Exception as e:
+            raise AssumeException(
+                f"Failed to load replay buffer from {path_to_load}: {e}"
+            )
+    else:
+        # create fresh buffer
+        buffer = ReplayBuffer(
             buffer_size=world.learning_role.learning_config.replay_buffer_size,
             obs_dim=world.learning_role.rl_algorithm.obs_dim,
             act_dim=world.learning_role.rl_algorithm.act_dim,
             n_rl_units=len(world.learning_role.rl_strats),
             device=world.learning_role.device,
             float_type=world.learning_role.float_type,
-        ),
+        )
+
+    inter_episodic_data = {
+        "buffer": buffer,
         "actors_and_critics": None,
         "max_eval": defaultdict(lambda: -1e9),
         "all_eval": defaultdict(list),
@@ -1296,6 +1345,29 @@ def run_learning(
             world.learning_role.rl_algorithm.save_params(
                 directory=f"{world.learning_role.learning_config.trained_policies_save_path}/last_policies"
             )
+            # also persist replay buffer alongside policies (configurable)
+            try:
+                if cfg_save_flag:
+                    save_path_cfg = (
+                        cfg_save_path
+                        if cfg_save_path is not None
+                        else default_buffer_path
+                    )
+                    # ensure directory exists will be handled by ReplayBuffer.save
+                    if (
+                        hasattr(world.learning_role, "buffer")
+                        and world.learning_role.buffer is not None
+                    ):
+                        logger.info(f"Saving replay buffer to {save_path_cfg}")
+                        world.learning_role.buffer.save(save_path_cfg)
+                        if os.path.exists(save_path_cfg):
+                            logger.info(f"Replay buffer saved: {save_path_cfg}")
+                        else:
+                            logger.warning(
+                                f"Replay buffer save attempted but file not found afterwards: {save_path_cfg}"
+                            )
+            except Exception:
+                logger.warning("Failed to save replay buffer")
 
     # container shutdown implicitly with new initialisation
     logger.info("################")

docs/source/learning_algorithm.rst

@@ -34,6 +34,10 @@ The following table shows the options that can be adjusted and gives a short exp
   continue_learning                       Whether to use pre-learned strategies and then continue learning. If True, loads existing policies from trained_policies_load_path and continues training. Note: Set True when you have a pretrained model and want incremental learning under new data or scenarios. Leave False for clean experiments. Default is False.
   trained_policies_save_path              The directory path - relative to the scenario's inputs_path - where newly trained RL policies (actor and critic networks) will be saved. Only needed when learning_mode is True. Value is set in setup_world(). Defaults otherwise to None.
   trained_policies_load_path              The directory path - relative to the scenario's inputs_path - from which pre-trained policies should be loaded. Needed when continue_learning is True or using pre-trained strategies. Default is None.
+    save_replay_buffer                      Whether ASSUME should persist the replay buffer when training finishes. Keeping this enabled is useful for long-running studies that may need to be resumed after an interruption. Default is True.
+    replay_buffer_save_path                 Optional path - relative to the scenario's inputs_path - where the replay buffer should be written. If omitted, ASSUME stores it next to the saved policies under ``last_policies/replay_buffer.npz``.
+    load_replay_buffer                      Whether an existing replay buffer should be restored before training starts. Use this together with ``continue_learning`` or interrupted-run recovery. Default is False.
+    replay_buffer_load_path                 Optional path - relative to the scenario's inputs_path - from which a replay buffer should be loaded. If omitted, ASSUME falls back to ``replay_buffer_save_path`` or the default location.
   min_bid_price                           The minimum bid price which limits the action of the actor to this price. Used to constrain the actor's output to a price range. Note: Best practice is to set this parameter as unconstraining as possible. When agent bid convergence is guaranteed to occur above zero, increasing the minimum bid value can reduce training times. Default is -100.0.
   max_bid_price                           The maximum bid price which limits the action of the actor to this price. Used to constrain the actor's output to a price range. Note: Align this with realistic market constraints. Too low = limited strategy space. Too high = noisy learning.  Default is 100.0.
   device                                  The device to use for PyTorch computations. Options include "cpu", "cuda", or specific CUDA devices like "cuda:0". Default is "cpu".
@@ -173,3 +177,28 @@ Yet, the buffer is quite large to store all observations also from multiple agen
 After a certain round of training runs which is defined in the config file the RL strategy is updated by calling the update function of the respective algorithms which calls the sample function of the replay buffer.
 The sample function returns a batch of experiences which is then used to update the RL strategy.
 For more information on the learning capabilities of ASSUME, see :doc:`learning`.
+
+
+Persisting and restoring replay buffers
+---------------------------------------
+
+ASSUME can optionally persist the replay buffer alongside trained policies so interrupted training runs can be resumed without recollecting the full experience set.
+
+The following learning-config items control this behaviour:
+
+- ``save_replay_buffer``: persist the replay buffer at the end of training.
+- ``replay_buffer_save_path``: optional explicit save path.
+- ``load_replay_buffer``: restore a previously saved buffer before training starts.
+- ``replay_buffer_load_path``: optional explicit load path.
+
+If no explicit save or load path is configured, ASSUME uses the default location ``last_policies/replay_buffer.npz`` inside the scenario inputs directory. A typical interrupted-run recovery setup looks like this:
+
+.. code-block:: yaml
+
+     learning_config:
+         continue_learning: true
+         trained_policies_load_path: learned_strategies/case_0_baseline/last_policies
+         save_replay_buffer: true
+         load_replay_buffer: true
+         replay_buffer_save_path: learned_strategies/case_0_baseline/last_policies/replay_buffer.npz
+         replay_buffer_load_path: learned_strategies/case_0_baseline/last_policies/replay_buffer.npz

docs/source/release_notes.rst

@@ -15,6 +15,7 @@ Upcoming Release
 
 **New Features:**
   - **Generic Forecasting Interface**: This interface enables to specify different forecast algorithms for preprocess, initialization and update during runtime. They can be specified in the config.yaml or unit csv files. For more information about currently implemented algorithms and how to specify them please read the documentation on Unit forecasts.
+  - **Replay-buffer persistence for RL**: Training runs can now save and optionally reload replay buffers, making it easier to resume long-running learning studies after crashes or scheduled interruptions.
 
 **Improvements:**
   - **In complex clearing, the solver instance is now created once during initialization of the clearing role and reused for each market clearing**. This improves performance for e.g. year-long simulations.

examples/inputs/example_02a/config.yaml

@@ -14,6 +14,10 @@ base:
     continue_learning: false
     trained_policies_save_path: null
     trained_policies_load_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001
@@ -59,6 +63,10 @@ base_lstm:
     learning_mode: true
     continue_learning: false
     trained_policies_save_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001
@@ -106,6 +114,10 @@ tiny:
     learning_mode: true
     continue_learning: false
     trained_policies_save_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001

examples/inputs/example_02b/config.yaml

@@ -14,6 +14,10 @@ base:
     continue_learning: false
     trained_policies_save_path: null
     trained_policies_load_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001
@@ -58,6 +62,10 @@ base_lstm:
     learning_mode: True
     continue_learning: False
     trained_policies_save_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     actor_architecture: lstm

examples/inputs/example_02c/config.yaml

@@ -14,6 +14,10 @@ base:
     continue_learning: false
     trained_policies_save_path: null
     trained_policies_load_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001

examples/inputs/example_02d/config.yaml

@@ -14,6 +14,10 @@ base:
     continue_learning: false
     trained_policies_save_path: null
     trained_policies_load_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001

examples/inputs/example_02e/config.yaml

@@ -14,6 +14,10 @@ base:
     continue_learning: false
     trained_policies_save_path: null
     trained_policies_load_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     actor_architecture: mlp
@@ -59,6 +63,10 @@ tiny:
     learning_mode: True
     continue_learning: False
     trained_policies_save_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 50
     algorithm: matd3
     actor_architecture: mlp

examples/inputs/example_03a/config.yaml

@@ -13,6 +13,10 @@ base_case_2019:
     learning_mode: true
     continue_learning: false
     trained_policies_save_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.0001

examples/inputs/example_03b/config.yaml

@@ -13,6 +13,10 @@ base_case_2021:
     learning_mode: True
     continue_learning: False
     trained_policies_save_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001

examples/inputs/example_03c/config.yaml

@@ -13,6 +13,10 @@ base_case_2019_with_storage:
     learning_mode: true
     continue_learning: false
     trained_policies_save_path: null
+    save_replay_buffer: true
+    replay_buffer_save_path: null
+    load_replay_buffer: false
+    replay_buffer_load_path: null
     max_bid_price: 100
     algorithm: matd3
     learning_rate: 0.001