Repository: AgileRL/AgileRL

Issue:
Bug with MultiDiscrete Observation Spaces
**Is your feature request related to a problem? Please describe.**
We already preprocess observations correctly for `MultiDiscrete` spaces there is just a bug because we don't include it in the `EvolvableAlgorithm.get_state_dim` method.

**Describe the solution you'd like**
Just need to add this line to the method
`elif isinstance(observation_space, spaces.MultiDiscrete):
            return (sum(observation_space.nvec),)`

Discussion:
@jaimesabalbermudez: Fixed in https://github.com/AgileRL/AgileRL/pull/325

Solution patch:
diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
index 90f873a..051cf84 100644
--- a/.github/workflows/python-app.yml
+++ b/.github/workflows/python-app.yml
@@ -28,6 +28,7 @@ jobs:
                   python -m pip install --upgrade pip
                   pip install ruff pytest pytest-cov black codespell yamlfmt isort pyupgrade
                   pip install -r requirements.txt
+                  pip install -e .
             - name: Run black
               run: |
                   black .
@@ -45,7 +46,7 @@ jobs:
                   pyupgrade --py38-plus
             - name: Test with pytest
               run: |
-                  pytest --cov=agilerl --cov-report=xml
+                  pytest --exitfirst --cov=agilerl --cov-report=xml
             - name: Upload coverage reports to Codecov
               uses: codecov/codecov-action@v3
               env:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 927f0a2..68dd31c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -32,14 +32,14 @@ repos:
                 - --skip=*.css,*.js,*.map,*.scss,*.svg
                 - --ignore-words-list=magent,pres,roate
     - repo: https://github.com/astral-sh/ruff-pre-commit
-      rev: v0.9.7
+      rev: v0.9.9
       hooks:
           - id: ruff
             args:
                 - --fix
                 - --ignore=E501
     - repo: https://github.com/PyCQA/isort
-      rev: 6.0.0
+      rev: 6.0.1
       hooks:
           - id: isort
             args: [--profile, black]
diff --git a/agilerl/algorithms/core/base.py b/agilerl/algorithms/core/base.py
index a5f0502..fdae0e2 100644
--- a/agilerl/algorithms/core/base.py
+++ b/agilerl/algorithms/core/base.py
@@ -1,6 +1,7 @@
 import copy
 import inspect
 from abc import ABC, ABCMeta, abstractmethod
+from importlib.metadata import version
 from typing import (
     Any,
     Callable,
@@ -158,6 +159,7 @@ def get_checkpoint_dict(agent: SelfEvolvableAlgorithm) -> Dict[str, Any]:
     network_info["network_names"] = network_attr_names
     network_info["optimizer_names"] = optimizer_attr_names
     attribute_dict["network_info"] = network_info
+    attribute_dict["agilerl_version"] = version("agilerl")
 
     attribute_dict.pop("accelerator", None)
     return attribute_dict
@@ -284,6 +286,8 @@ class EvolvableAlgorithm(ABC, metaclass=RegistryMeta):
             )
         elif isinstance(observation_space, spaces.Discrete):
             return (observation_space.n,)
+        elif isinstance(observation_space, spaces.MultiDiscrete):
+            return (sum(observation_space.nvec),)
         elif isinstance(observation_space, spaces.Box):
             return observation_space.shape
         elif isinstance(observation_space, spaces.Dict):
@@ -848,6 +852,7 @@ class EvolvableAlgorithm(ABC, metaclass=RegistryMeta):
         )
 
         # Reconstruct evolvable modules in algorithm
+        print("Checkpoint: ", checkpoint)
         network_info: Optional[Dict[str, Dict[str, Any]]] = checkpoint.get(
             "network_info"
         )
@@ -913,6 +918,8 @@ class EvolvableAlgorithm(ABC, metaclass=RegistryMeta):
         checkpoint["accelerator"] = accelerator
         checkpoint["device"] = device
         self = cls(**class_init_dict)
+        registry: MutationRegistry = checkpoint["registry"]
+        self.registry = registry
 
         # Reconstruct optimizers in algorithm
         optimizer_names = network_info["optimizer_names"]
@@ -959,13 +966,6 @@ class EvolvableAlgorithm(ABC, metaclass=RegistryMeta):
         for name, optimizer in loaded_optimizers.items():
             setattr(self, name, optimizer)
 
-        # Check loaded registry is consistent with the algorithm
-        if checkpoint["registry"] != self.registry:
-            raise ValueError(
-                "Loaded registry does not match the algorithm's registry. Please make "
-                "sure you are loading the checkpoint with the correct algorithm."
-            )
-
         # Assign other attributes to the algorithm
         for attribute in EvolvableAlgorithm.inspect_attributes(self).keys():
             setattr(self, attribute, checkpoint[attribute])
@@ -1098,8 +1098,8 @@ class MultiAgentRLAlgorithm(EvolvableAlgorithm, ABC):
 
     :param observation_spaces: The observation spaces of the agent environments.
     :type observation_spaces: List[spaces.Space]
-    :param action_space: The action spaces of the agent environments.
-    :type action_space: List[spaces.Space]
+    :param action_spaces: The action spaces of the agent environments.
+    :type action_spaces: List[spaces.Space]
     :param agent_ids: The agent IDs of the agents in the environment.
     :type agent_ids: List[int]
     :param index: The index of the individual in the population.
@@ -1169,8 +1169,7 @@ class MultiAgentRLAlgorithm(EvolvableAlgorithm, ABC):
             isinstance(space, spaces.Discrete) for space in observation_spaces
         )
         self.discrete_actions = all(
-            isinstance(space, (spaces.Discrete, spaces.MultiDiscrete))
-            for space in action_spaces
+            isinstance(space, spaces.Discrete) for space in action_spaces
         )
 
         # For continuous action spaces, store the min and max action values
diff --git a/agilerl/components/multi_agent_replay_buffer.py b/agilerl/components/multi_agent_replay_buffer.py
index c6de4e1..90e162b 100644
--- a/agilerl/components/multi_agent_replay_buffer.py
+++ b/agilerl/components/multi_agent_replay_buffer.py
@@ -148,7 +148,6 @@ class MultiAgentReplayBuffer:
                     ts = ts.astype(np.uint8)
 
                 if not np_array:
-                    print(ts)
                     ts = obs_to_tensor(ts, self.device)
 
                 transition[field][agent_id] = ts
diff --git a/docs/distributed_training/index.rst b/docs/distributed_training/index.rst
index 4a15575..f3fe057 100644
--- a/docs/distributed_training/index.rst
+++ b/docs/distributed_training/index.rst
@@ -61,7 +61,7 @@ Example distributed training loop:
     num_envs = 8
     env = make_vect_envs("LunarLander-v2", num_envs=num_envs)  # Create environment
     observation_space = env.single_observation_space
-    action_space - env.single_action_space
+    action_space = env.single_action_space
 
     if INIT_HP['CHANNELS_LAST']:
         observation_space = observation_space_channels_to_first(observation_space)
diff --git a/docs/off_policy/index.rst b/docs/off_policy/index.rst
index 1e1726a..5f88601 100644
--- a/docs/off_policy/index.rst
+++ b/docs/off_policy/index.rst
@@ -71,7 +71,7 @@ are more likely to remain present in the population. The sequence of evolution (
     env = make_vect_envs("LunarLander-v2", num_envs=num_envs)  # Create environment
 
     observation_space = env.single_observation_space
-    action_space - env.single_action_space
+    action_space = env.single_action_space
 
     if INIT_HP['CHANNELS_LAST']:
         observation_space = observation_space_channels_to_first(observation_space)
@@ -186,7 +186,7 @@ Alternatively, use a custom training loop. Combining all of the above:
     env = make_vect_envs("LunarLander-v2", num_envs=num_envs)  # Create environment
 
     observation_space = env.single_observation_space
-    action_space - env.single_action_space
+    action_space = env.single_action_space
 
     if INIT_HP['CHANNELS_LAST']:
         observation_space = observation_space_channels_to_first(observation_space)
diff --git a/docs/offline_training/index.rst b/docs/offline_training/index.rst
index 673c49f..d058c36 100644
--- a/docs/offline_training/index.rst
+++ b/docs/offline_training/index.rst
@@ -53,7 +53,7 @@ are more likely to remain present in the population. The sequence of evolution (
     dataset = h5py.File("data/cartpole/cartpole_random_v1.1.0.h5", "r")  # Load dataset
 
     observation_space = env.single_observation_space
-    action_space - env.single_action_space
+    action_space = env.single_action_space
     if INIT_HP['CHANNELS_LAST']:
         observation_space = observation_space_channels_to_first(observation_space)
 
@@ -191,7 +191,7 @@ Alternatively, use a custom training loop. Combining all of the above:
     dataset = h5py.File("data/cartpole/cartpole_random_v1.1.0.h5", "r")  # Load dataset
 
     observation_space = env.single_observation_space
-    action_space - env.single_action_space
+    action_space = env.single_action_space
     if INIT_HP['CHANNELS_LAST']:
         observation_space = observation_space_channels_to_first(observation_space)
 
diff --git a/docs/on_policy/index.rst b/docs/on_policy/index.rst
index 3b9c176..9a60216 100644
--- a/docs/on_policy/index.rst
+++ b/docs/on_policy/index.rst
@@ -69,7 +69,7 @@ are more likely to remain present in the population. The sequence of evolution (
     env = make_vect_envs("LunarLander-v2", num_envs=num_envs)  # Create environment
 
     observation_space = env.single_observation_space
-    action_space - env.single_action_space
+    action_space = env.single_action_space
     if INIT_HP['CHANNELS_LAST']:
         observation_space = observation_space_channels_to_first(observation_space)
 
@@ -166,7 +166,7 @@ Alternatively, use a custom on-policy training loop:
     env = make_vect_envs("LunarLander-v2", num_envs=num_envs)  # Create environment
 
     observation_space = env.single_observation_space
-    action_space - env.single_action_space
+    action_space = env.single_action_space
     if INIT_HP['CHANNELS_LAST']:
         observation_space = observation_space_channels_to_first(observation_space)
 
diff --git a/pyproject.toml b/pyproject.toml
index 5008109..d5f8aa3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agilerl"
-version = "2.0.5"
+version = "2.0.6"
 
 description = "AgileRL is a deep reinforcement learning library focused on improving RL development through RLOps."
 authors = ["Nick Ustaran-Anderegg <dev@agilerl.com>"]
@@ -18,7 +18,7 @@ gymnasium = "^0.29.1"
 google-cloud-storage = "^2.5.0"
 h5py = "^3.8.0"
 hydra-core = "^1.3.2"
-matplotlib = "~3.4.3"
+matplotlib = "~3.9.4"
 minari = "^0.5.2"
 numpy = "^1.26.4"
 omegaconf = "^2.3.0"
@@ -32,7 +32,7 @@ redis = "^4.4.4"
 scipy = "^1.12.0"
 SuperSuit = "^3.9.0"
 termcolor = "^1.1.0"
-torch = "^2.5.1"
+torch = "2.5.1"
 tensordict = "^0.6"
 tqdm = "^4.66.4"
 transformers = "^4.41.2"
diff --git a/requirements.txt b/requirements.txt
index e40e418..486196e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,7 +7,7 @@ gymnasium==0.29.1
 h5py==3.8.0
 hydra-core==1.3.2
 jax[cpu]==0.4.31
-matplotlib==3.4.3
+matplotlib==3.9.4
 minari[all]==0.5.2
 numpy==1.26.4
 omegaconf==2.3.0
diff --git a/tutorials/PettingZoo/agilerl_maddpg.py b/tutorials/PettingZoo/agilerl_maddpg.py
index fdcf77d..ce1ec12 100644
--- a/tutorials/PettingZoo/agilerl_maddpg.py
+++ b/tutorials/PettingZoo/agilerl_maddpg.py
@@ -13,6 +13,7 @@ from pettingzoo.atari import space_invaders_v2
 from tqdm import trange
 
 from agilerl.algorithms.core.registry import HyperparameterConfig, RLParameter
+from agilerl.algorithms.maddpg import MADDPG
 from agilerl.components.multi_agent_replay_buffer import MultiAgentReplayBuffer
 from agilerl.utils.algo_utils import obs_channels_to_first
 from agilerl.utils.utils import create_population, observation_space_channels_to_first
@@ -87,7 +88,7 @@ if __name__ == "__main__":
     )
 
     # Create a population ready for evolutionary hyper-parameter optimisation
-    agent = create_population(
+    agent: MADDPG = create_population(
         INIT_HP["ALGO"],
         observation_spaces,
         action_spaces,
@@ -100,7 +101,7 @@ if __name__ == "__main__":
     )[0]
 
     # Configure the multi-agent replay buffer
-    field_names = ["state", "action", "reward", "next_state", "done"]
+    field_names = ["obs", "action", "reward", "next_obs", "done"]
     memory = MultiAgentReplayBuffer(
         INIT_HP["MEMORY_SIZE"],
         field_names=field_names,
@@ -122,20 +123,18 @@ if __name__ == "__main__":
     print("Training...")
     pbar = trange(max_steps, unit="step")
     while np.less(agent.steps[-1], max_steps):
-        state, info = env.reset()  # Reset environment at start of episode
+        obs, info = env.reset()  # Reset environment at start of episode
         scores = np.zeros((num_envs, len(agent_ids)))
         completed_episode_scores = []
         steps = 0
         if INIT_HP["CHANNELS_LAST"]:
-            state = {
-                agent_id: obs_channels_to_first(s) for agent_id, s in state.items()
-            }
+            obs = {agent_id: obs_channels_to_first(s) for agent_id, s in obs.items()}
 
         for idx_step in range(training_steps // num_envs):
 
             # Get next action from agent
             cont_actions, discrete_action = agent.get_action(
-                states=state, training=True, infos=info
+                obs=obs, training=True, infos=info
             )
             if agent.discrete_actions:
                 action = discrete_action
@@ -144,7 +143,7 @@ if __name__ == "__main__":
 
             # Act in environment
             action = {agent: env.action_space(agent).sample() for agent in env.agents}
-            next_state, reward, termination, truncation, info = env.step(action)
+            next_obs, reward, termination, truncation, info = env.step(action)
             if not termination:
                 assert False
             scores += np.array(list(reward.values())).transpose()
@@ -153,17 +152,17 @@ if __name__ == "__main__":
 
             # Image processing if necessary for the environment
             if INIT_HP["CHANNELS_LAST"]:
-                next_state = {
+                next_obs = {
                     agent_id: obs_channels_to_first(ns)
-                    for agent_id, ns in next_state.items()
+                    for agent_id, ns in next_obs.items()
                 }
 
             # Save experiences to replay buffer
             memory.save_to_memory(
-                state,
+                obs,
                 cont_actions,
                 reward,
-                next_state,
+                next_obs,
                 termination,
                 is_vectorised=True,
             )
@@ -181,6 +180,7 @@ if __name__ == "__main__":
                     experiences = memory.sample(agent.batch_size)
                     # Learn according to agent's RL algorithm
                     agent.learn(experiences)
+
             # Handle num_envs > learn step; learn multiple times per step in env
             elif len(memory) >= agent.batch_size and memory.counter > learning_delay:
                 for _ in range(num_envs // agent.learn_step):
@@ -189,7 +189,7 @@ if __name__ == "__main__":
                     # Learn according to agent's RL algorithm
                     agent.learn(experiences)
 
-            state = next_state
+            obs = next_obs
 
             # Calculate scores and reset noise for finished episodes
             reset_noise_indices = []
diff --git a/tutorials/PettingZoo/agilerl_matd3.py b/tutorials/PettingZoo/agilerl_matd3.py
index 7bb0222..59b1e4f 100644
--- a/tutorials/PettingZoo/agilerl_matd3.py
+++ b/tutorials/PettingZoo/agilerl_matd3.py
@@ -4,12 +4,14 @@ Authors: Michael (https://github.com/mikepratt1), Nickua (https://github.com/nic
 """
 
 import os
+from typing import List
 
 import numpy as np
 import torch
 from pettingzoo.mpe import simple_speaker_listener_v4
 from tqdm import trange
 
+from agilerl.algorithms import MATD3
 from agilerl.algorithms.core.registry import HyperparameterConfig, RLParameter
 from agilerl.components.multi_agent_replay_buffer import MultiAgentReplayBuffer
 from agilerl.hpo.mutation import Mutations
@@ -79,7 +81,7 @@ if __name__ == "__main__":
     )
 
     # Create a population ready for evolutionary hyper-parameter optimisation
-    pop = create_population(
+    pop: List[MATD3] = create_population(
         INIT_HP["ALGO"],
         observation_spaces,
         action_spaces,
@@ -92,7 +94,7 @@ if __name__ == "__main__":
     )
 
     # Configure the multi-agent replay buffer
-    field_names = ["state", "action", "reward", "next_state", "done"]
+    field_names = ["obs", "action", "reward", "next_obs", "done"]
     memory = MultiAgentReplayBuffer(
         INIT_HP["MEMORY_SIZE"],
         field_names=field_names,
@@ -137,20 +139,20 @@ if __name__ == "__main__":
     while np.less([agent.steps[-1] for agent in pop], max_steps).all():
         pop_episode_scores = []
         for agent in pop:  # Loop through population
-            state, info = env.reset()  # Reset environment at start of episode
+            obs, info = env.reset()  # Reset environment at start of episode
             scores = np.zeros(num_envs)
             completed_episode_scores = []
             steps = 0
             if INIT_HP["CHANNELS_LAST"]:
-                state = {
-                    agent_id: obs_channels_to_first(s) for agent_id, s in state.items()
+                obs = {
+                    agent_id: obs_channels_to_first(s) for agent_id, s in obs.items()
                 }
 
             for idx_step in range(evo_steps // num_envs):
 
                 # Get next action from agent
                 cont_actions, discrete_action = agent.get_action(
-                    states=state, training=True, infos=info
+                    obs=obs, training=True, infos=info
                 )
                 if agent.discrete_actions:
                     action = discrete_action
@@ -158,7 +160,7 @@ if __name__ == "__main__":
                     action = cont_actions
 
                 # Act in environment
-                next_state, reward, termination, truncation, info = env.step(action)
+                next_obs, reward, termination, truncation, info = env.step(action)
 
                 scores += np.sum(np.array(list(reward.values())).transpose(), axis=-1)
                 total_steps += num_envs
@@ -166,17 +168,17 @@ if __name__ == "__main__":
 
                 # Image processing if necessary for the environment
                 if INIT_HP["CHANNELS_LAST"]:
-                    next_state = {
+                    next_obs = {
                         agent_id: obs_channels_to_first(ns)
-                        for agent_id, ns in next_state.items()
+                        for agent_id, ns in next_obs.items()
                     }
 
                 # Save experiences to replay buffer
                 memory.save_to_memory(
-                    state,
+                    obs,
                     cont_actions,
                     reward,
-                    next_state,
+                    next_obs,
                     termination,
                     is_vectorised=True,
                 )
@@ -204,7 +206,7 @@ if __name__ == "__main__":
                         # Learn according to agent's RL algorithm
                         agent.learn(experiences)
 
-                state = next_state
+                obs = next_obs
 
                 # Calculate scores and reset noise for finished episodes
                 reset_noise_indices = []


Test patch:
diff --git a/tests/test_algorithms/test_base.py b/tests/test_algorithms/test_base.py
index f1a048d..0e793a6 100644
--- a/tests/test_algorithms/test_base.py
+++ b/tests/test_algorithms/test_base.py
@@ -1,4 +1,7 @@
+from pathlib import Path
+
 import pytest
+import torch
 import torch.optim as optim
 from gymnasium import spaces
 
@@ -16,6 +19,7 @@ from tests.helper_functions import (
     generate_discrete_space,
     generate_multi_agent_box_spaces,
     generate_multi_agent_discrete_spaces,
+    generate_multidiscrete_space,
     generate_random_box_space,
 )
 
@@ -36,8 +40,8 @@ def multi_input_config():
 
 
 class DummyRLAlgorithm(RLAlgorithm):
-    def __init__(self, *args, lr=True, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, observation_space, action_space, index, lr=True, **kwargs):
+        super().__init__(observation_space, action_space, index, **kwargs)
 
         num_outputs = (
             self.action_space.n
@@ -59,6 +63,10 @@ class DummyRLAlgorithm(RLAlgorithm):
                 else self.observation_space.n
             )
             self.dummy_actor = EvolvableMLP(num_inputs, num_outputs, hidden_size=[8])
+        elif isinstance(self.observation_space, spaces.MultiDiscrete):
+            # Handle MultiDiscrete spaces
+            num_inputs = len(self.observation_space.nvec)
+            self.dummy_actor = EvolvableMLP(num_inputs, num_outputs, hidden_size=[8])
         elif isinstance(self.observation_space, (spaces.Dict, spaces.Tuple)):
             self.dummy_actor = EvolvableMultiInput(
                 self.observation_space,
@@ -85,8 +93,8 @@ class DummyRLAlgorithm(RLAlgorithm):
 
 
 class DummyMARLAlgorithm(MultiAgentRLAlgorithm):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, observation_spaces, action_spaces, agent_ids, index, **kwargs):
+        super().__init__(observation_spaces, action_spaces, agent_ids, index, **kwargs)
 
         def create_actor(idx):
             obs_space = self.observation_spaces[idx]
@@ -221,3 +229,291 @@ def test_incorrect_hp_config():
             index=0,
             hp_config=hp_config,
         )
+
+
+@pytest.mark.parametrize(
+    "with_hp_config",
+    [
+        False,
+        True,
+    ],
+)
+@pytest.mark.parametrize(
+    "observation_space",
+    [
+        generate_random_box_space((4,)),
+        generate_discrete_space(4),
+        generate_dict_or_tuple_space(1, 1, dict_space=True),
+        generate_dict_or_tuple_space(1, 1, dict_space=False),
+        generate_multidiscrete_space(2, 2),
+    ],
+)
+def test_save_load_checkpoint_single_agent(tmpdir, with_hp_config, observation_space):
+    action_space = generate_discrete_space(4)
+    # Initialize the dummy agent
+    hp_config = None
+    if with_hp_config:
+        hp_config = HyperparameterConfig(lr=RLParameter(min=0.05, max=0.2))
+        agent = DummyRLAlgorithm(
+            observation_space, action_space, index=0, hp_config=hp_config
+        )
+    else:
+        agent = DummyRLAlgorithm(observation_space, action_space, index=0)
+
+    # Save the checkpoint to a file
+    checkpoint_path = Path(tmpdir) / "checkpoint.pth"
+    agent.save_checkpoint(checkpoint_path)
+
+    # Load the saved checkpoint file
+    checkpoint = torch.load(checkpoint_path)
+
+    # Check if the loaded checkpoint has the correct keys
+    assert "dummy_actor_init_dict" in checkpoint["network_info"]["modules"]
+    assert "dummy_actor_state_dict" in checkpoint["network_info"]["modules"]
+    assert "dummy_optimizer_state_dict" in checkpoint["network_info"]["optimizers"]
+    assert "lr" in checkpoint
+    assert "index" in checkpoint
+    assert "scores" in checkpoint
+    assert "fitness" in checkpoint
+    assert "steps" in checkpoint
+
+    # Create a new agent with the same hp_config if needed
+    new_agent = DummyRLAlgorithm(
+        observation_space,
+        action_space,
+        index=1,  # Different index to verify it gets overwritten
+        hp_config=hp_config,
+    )
+
+    # Load checkpoint
+    new_agent.load_checkpoint(checkpoint_path)
+
+    # Check if properties and weights are loaded correctly
+    assert isinstance(
+        new_agent.dummy_actor, (EvolvableMLP, EvolvableCNN, EvolvableMultiInput)
+    )
+    assert new_agent.lr == agent.lr
+    assert str(new_agent.dummy_actor.state_dict()) == str(
+        agent.dummy_actor.state_dict()
+    )
+    assert new_agent.index == agent.index
+    assert new_agent.scores == agent.scores
+    assert new_agent.fitness == agent.fitness
+    assert new_agent.steps == agent.steps
+
+
+@pytest.mark.parametrize(
+    "with_hp_config",
+    [
+        False,
+        True,
+    ],
+)
+@pytest.mark.parametrize(
+    "observation_spaces",
+    [
+        generate_multi_agent_box_spaces(2, (4,)),
+        generate_multi_agent_discrete_spaces(2, 4),
+        gen_multi_agent_dict_or_tuple_spaces(2, 1, 1, dict_space=True),
+        gen_multi_agent_dict_or_tuple_spaces(2, 1, 1, dict_space=False),
+    ],
+)
+def test_save_load_checkpoint_multi_agent(tmpdir, with_hp_config, observation_spaces):
+    # Initialize the dummy multi-agent
+    agent_ids = ["agent1", "agent2"]
+    action_spaces = generate_multi_agent_discrete_spaces(2, 4)
+
+    hp_config = None
+    if with_hp_config:
+        hp_config = HyperparameterConfig(lr=RLParameter(min=0.05, max=0.2))
+        agent = DummyMARLAlgorithm(
+            observation_spaces,
+            action_spaces,
+            agent_ids=agent_ids,
+            index=0,
+            hp_config=hp_config,
+        )
+    else:
+        agent = DummyMARLAlgorithm(
+            observation_spaces, action_spaces, agent_ids=agent_ids, index=0
+        )
+
+    # Save the checkpoint to a file
+    checkpoint_path = Path(tmpdir) / "checkpoint.pth"
+    agent.save_checkpoint(checkpoint_path)
+
+    # Load the saved checkpoint file
+    checkpoint = torch.load(checkpoint_path)
+
+    # Check if the loaded checkpoint has the correct keys
+    assert "dummy_actors_init_dict" in checkpoint["network_info"]["modules"]
+    assert "dummy_actors_state_dict" in checkpoint["network_info"]["modules"]
+    assert "dummy_optimizer_state_dict" in checkpoint["network_info"]["optimizers"]
+    assert "lr" in checkpoint
+    assert "index" in checkpoint
+    assert "scores" in checkpoint
+    assert "fitness" in checkpoint
+    assert "steps" in checkpoint
+    assert "agent_ids" in checkpoint
+
+    # Create a new agent with the same hp_config if needed
+    new_agent = DummyMARLAlgorithm(
+        observation_spaces,
+        action_spaces,
+        agent_ids=agent_ids,
+        index=1,  # Different index to verify it gets overwritten
+        hp_config=hp_config,
+    )
+
+    # Load checkpoint
+    new_agent.load_checkpoint(checkpoint_path)
+
+    # Check if properties and weights are loaded correctly
+    for i in range(len(agent.dummy_actors)):
+        assert isinstance(
+            new_agent.dummy_actors[i], (EvolvableMLP, EvolvableCNN, EvolvableMultiInput)
+        )
+        assert str(new_agent.dummy_actors[i].state_dict()) == str(
+            agent.dummy_actors[i].state_dict()
+        )
+
+    assert new_agent.lr == agent.lr
+    assert new_agent.index == agent.index
+    assert new_agent.scores == agent.scores
+    assert new_agent.fitness == agent.fitness
+    assert new_agent.steps == agent.steps
+    assert new_agent.agent_ids == agent.agent_ids
+
+
+@pytest.mark.parametrize(
+    "device, with_hp_config",
+    [
+        ("cpu", False),
+        ("cpu", True),
+    ],
+)
+@pytest.mark.parametrize(
+    "observation_space",
+    [
+        generate_random_box_space((4,)),
+        generate_discrete_space(4),
+        generate_dict_or_tuple_space(1, 1, dict_space=True),
+        generate_dict_or_tuple_space(1, 1, dict_space=False),
+        generate_multidiscrete_space(2, 2),
+    ],
+)
+@pytest.mark.parametrize(
+    "action_space",
+    [
+        generate_random_box_space((2,)),
+        generate_discrete_space(4),
+    ],
+)
+def test_load_from_pretrained_single_agent(
+    device, tmpdir, with_hp_config, observation_space, action_space
+):
+    # Initialize the dummy agent
+    hp_config = None
+    if with_hp_config:
+        hp_config = HyperparameterConfig(lr=RLParameter(min=0.05, max=0.2))
+        agent = DummyRLAlgorithm(
+            observation_space, action_space, index=0, hp_config=hp_config
+        )
+    else:
+        agent = DummyRLAlgorithm(observation_space, action_space, index=0)
+
+    # Save the checkpoint to a file
+    checkpoint_path = Path(tmpdir) / "checkpoint.pth"
+    agent.save_checkpoint(checkpoint_path)
+
+    # Create new agent object using the class method
+    new_agent = DummyRLAlgorithm.load(checkpoint_path, device=device)
+
+    # Check if properties and weights are loaded correctly
+    assert new_agent.observation_space == agent.observation_space
+    assert new_agent.action_space == agent.action_space
+    assert isinstance(
+        new_agent.dummy_actor, (EvolvableMLP, EvolvableCNN, EvolvableMultiInput)
+    )
+    assert new_agent.lr == agent.lr
+    assert str(new_agent.dummy_actor.to("cpu").state_dict()) == str(
+        agent.dummy_actor.state_dict()
+    )
+    assert new_agent.index == agent.index
+    assert new_agent.scores == agent.scores
+    assert new_agent.fitness == agent.fitness
+    assert new_agent.steps == agent.steps
+
+
+@pytest.mark.parametrize(
+    "device, with_hp_config",
+    [
+        ("cpu", False),
+        ("cpu", True),
+    ],
+)
+@pytest.mark.parametrize(
+    "observation_spaces",
+    [
+        generate_multi_agent_box_spaces(2, (4,)),
+        generate_multi_agent_discrete_spaces(2, 4),
+        gen_multi_agent_dict_or_tuple_spaces(2, 1, 1, dict_space=True),
+        gen_multi_agent_dict_or_tuple_spaces(2, 1, 1, dict_space=False),
+    ],
+)
+@pytest.mark.parametrize(
+    "action_spaces",
+    [
+        generate_multi_agent_box_spaces(2, (2,)),
+        generate_multi_agent_discrete_spaces(2, 4),
+    ],
+)
+def test_load_from_pretrained_multi_agent(
+    device, tmpdir, with_hp_config, observation_spaces, action_spaces
+):
+    # Initialize the dummy multi-agent
+    agent_ids = ["agent1", "agent2"]
+
+    hp_config = None
+    if with_hp_config:
+        hp_config = HyperparameterConfig(lr=RLParameter(min=0.05, max=0.2))
+        agent = DummyMARLAlgorithm(
+            observation_spaces,
+            action_spaces,
+            agent_ids=agent_ids,
+            index=0,
+            hp_config=hp_config,
+        )
+    else:
+        agent = DummyMARLAlgorithm(
+            observation_spaces, action_spaces, agent_ids=agent_ids, index=0
+        )
+
+    # Save the checkpoint to a file
+    checkpoint_path = Path(tmpdir) / "checkpoint.pth"
+    agent.save_checkpoint(checkpoint_path)
+
+    # Create new agent object using the class method
+    new_agent = DummyMARLAlgorithm.load(checkpoint_path, device=device)
+
+    # Check if properties and weights are loaded correctly
+    for i, agent_id in enumerate(agent_ids):
+        assert (
+            new_agent.observation_space[agent_id] == agent.observation_space[agent_id]
+        )
+        assert new_agent.action_space[agent_id] == agent.action_space[agent_id]
+
+    for i in range(len(agent.dummy_actors)):
+        assert isinstance(
+            new_agent.dummy_actors[i], (EvolvableMLP, EvolvableCNN, EvolvableMultiInput)
+        )
+        assert str(new_agent.dummy_actors[i].to("cpu").state_dict()) == str(
+            agent.dummy_actors[i].state_dict()
+        )
+
+    assert new_agent.lr == agent.lr
+    assert new_agent.index == agent.index
+    assert new_agent.scores == agent.scores
+    assert new_agent.fitness == agent.fitness
+    assert new_agent.steps == agent.steps
+    assert new_agent.agent_ids == agent.agent_ids
diff --git a/tests/test_algorithms/test_cqn.py b/tests/test_algorithms/test_cqn.py
index 30c594c..4a937c9 100644
--- a/tests/test_algorithms/test_cqn.py
+++ b/tests/test_algorithms/test_cqn.py
@@ -825,6 +825,9 @@ def test_load_from_pretrained(device, accelerator, tmpdir):
     checkpoint_path = Path(tmpdir) / "checkpoint.pth"
     cqn.save_checkpoint(checkpoint_path)
 
+    checkpoint = torch.load(checkpoint_path, map_location=device, pickle_module=dill)
+    assert "agilerl_version" in checkpoint
+
     # Create new agent object
     new_cqn = CQN.load(checkpoint_path, device=device, accelerator=accelerator)
 
diff --git a/tests/test_algorithms/test_ddpg.py b/tests/test_algorithms/test_ddpg.py
index 5cb1ca8..e46e802 100644
--- a/tests/test_algorithms/test_ddpg.py
+++ b/tests/test_algorithms/test_ddpg.py
@@ -1212,6 +1212,9 @@ def test_load_from_pretrained(device, accelerator, tmpdir):
     checkpoint_path = Path(tmpdir) / "checkpoint.pth"
     ddpg.save_checkpoint(checkpoint_path)
 
+    checkpoint = torch.load(checkpoint_path, map_location=device, pickle_module=dill)
+    assert "agilerl_version" in checkpoint
+
     # Create new agent object
     new_ddpg = DDPG.load(checkpoint_path, device=device, accelerator=accelerator)