[RLlib] Algorithm Level Checkpointing with Learner and RL Modules (ra…

…y-project#34717)

Signed-off-by: Avnish <[email protected]>

rllib/BUILD

@@ -1982,6 +1982,13 @@ py_test(
  srcs = ["core/learner/torch/tests/test_torch_learner.py"]
 )
 
+py_test(
+ name ="tests/test_algorithm_save_load_checkpoint_learner",
+ tags = ["team:rllib", "core"],
+ size = "medium",
+ srcs = ["tests/test_algorithm_save_load_checkpoint_learner.py"]
+)
+
 py_test(
  name = "test_bc_algorithm",
  tags = ["team:rllib", "core"],

rllib/algorithms/algorithm.py

@@ -76,6 +76,7 @@
 )
 from ray.rllib.utils.checkpoints import (
  CHECKPOINT_VERSION,
+ CHECKPOINT_VERSION_LEARNER,
  get_checkpoint_info,
  try_import_msgpack,
 )
@@ -2077,6 +2078,14 @@ def save_checkpoint(self, checkpoint_dir: str) -> str:
  policy_state.pkl
  pol_2/
  policy_state.pkl
+ learner/
+ learner_state.json
+ module_state/
+ module_1/
+ ...
+ optimizer_state/
+ optimizers_module_1/
+ ...
  rllib_checkpoint.json
  algorithm_state.pkl
 
@@ -2099,7 +2108,10 @@ def save_checkpoint(self, checkpoint_dir: str) -> str:
  policy_states = state["worker"].pop("policy_states", {})
 
  # Add RLlib checkpoint version.
- state["checkpoint_version"] = CHECKPOINT_VERSION
+ if self.config._enable_learner_api:
+ state["checkpoint_version"] = CHECKPOINT_VERSION_LEARNER
+ else:
+ state["checkpoint_version"] = CHECKPOINT_VERSION
 
  # Write state (w/o policies) to disk.
  state_file = os.path.join(checkpoint_dir, "algorithm_state.pkl")
@@ -2130,21 +2142,24 @@ def save_checkpoint(self, checkpoint_dir: str) -> str:
  policy = self.get_policy(pid)
  policy.export_checkpoint(policy_dir, policy_state=policy_state)
 
+ # if we are using the learner API, save the learner group state
+ if self.config._enable_learner_api:
+ learner_state_dir = os.path.join(checkpoint_dir, "learner")
+ self.learner_group.save_state(learner_state_dir)
+
  return checkpoint_dir
 
  @override(Trainable)
- def load_checkpoint(self, checkpoint: Union[Dict, str]) -> None:
+ def load_checkpoint(self, checkpoint: str) -> None:
  # Checkpoint is provided as a directory name.
  # Restore from the checkpoint file or dir.
- if isinstance(checkpoint, str):
- checkpoint_info = get_checkpoint_info(checkpoint)
- checkpoint_data = Algorithm._checkpoint_info_to_algorithm_state(
- checkpoint_info
- )
- # Checkpoint is a checkpoint-as-dict -> Restore state from it as-is.
- else:
- checkpoint_data = checkpoint
+
+ checkpoint_info = get_checkpoint_info(checkpoint)
+ checkpoint_data = Algorithm._checkpoint_info_to_algorithm_state(checkpoint_info)
  self.__setstate__(checkpoint_data)
+ if self.config._enable_learner_api:
+ learner_state_dir = os.path.join(checkpoint, "learner")
+ self.learner_group.load_state(learner_state_dir)
 
  @override(Trainable)
  def log_result(self, result: ResultDict) -> None:

rllib/algorithms/ppo/tests/test_ppo_learner.py

@@ -2,6 +2,7 @@
 import unittest
 import numpy as np
 import torch
+import tempfile
 import tensorflow as tf
 import tree # pip install dm-tree
 
@@ -74,8 +75,8 @@ def test_loss(self):
  )
 
  for fw in framework_iterator(config, ("tf2", "torch"), with_eager_tracing=True):
- trainer = config.build()
- policy = trainer.get_policy()
+ algo = config.build()
+ policy = algo.get_policy()
 
  train_batch = SampleBatch(FAKE_BATCH)
  train_batch = compute_gae_for_sample_batch(policy, train_batch)
@@ -109,14 +110,58 @@ def test_loss(self):
  )
  learner_group = learner_group_config.build()
 
- # load the trainer weights onto the learner_group
- learner_group.set_weights(trainer.get_weights())
+ # load the algo weights onto the learner_group
+ learner_group.set_weights(algo.get_weights())
  results = learner_group.update(train_batch.as_multi_agent())
 
  learner_group_loss = results[ALL_MODULES]["total_loss"]
 
  check(learner_group_loss, policy_loss)
 
+ def test_save_load_state(self):
+ """Tests saving and loading the state of the PPO Learner Group."""
+ config = (
+ ppo.PPOConfig()
+ .environment("CartPole-v1")
+ .rollouts(
+ num_rollout_workers=0,
+ )
+ .training(
+ gamma=0.99,
+ model=dict(
+ fcnet_hiddens=[10, 10],
+ fcnet_activation="linear",
+ vf_share_layers=False,
+ ),
+ _enable_learner_api=True,
+ )
+ .rl_module(
+ _enable_rl_module_api=True,
+ )
+ )
+ algo = config.build()
+ policy = algo.get_policy()
+
+ for fw in framework_iterator(config, ("tf2", "torch"), with_eager_tracing=True):
+ algo_config = config.copy(copy_frozen=False)
+ algo_config.validate()
+ algo_config.freeze()
+ learner_group_config = algo_config.get_learner_group_config(
+ SingleAgentRLModuleSpec(
+ module_class=algo_config.rl_module_spec.module_class,
+ observation_space=policy.observation_space,
+ action_space=policy.action_space,
+ model_config_dict=policy.config["model"],
+ catalog_class=PPOCatalog,
+ )
+ )
+ learner_group1 = learner_group_config.build()
+ learner_group2 = learner_group_config.build()
+ with tempfile.TemporaryDirectory() as tmpdir:
+ learner_group1.save_state(tmpdir)
+ learner_group2.load_state(tmpdir)
+ check(learner_group1.get_state(), learner_group2.get_state())
+
 
 if __name__ == "__main__":
  import pytest

rllib/core/learner/learner.py

@@ -906,6 +906,16 @@ def save_state(self, path: Union[str, pathlib.Path]) -> None:
  NOTE: if path doesn't exist, then a new directory will be created. otherwise, it
  will be appended to.
 
+ the state of the learner is saved in the following format:
+
+ checkpoint_dir/
+ learner_state.json
+ module_state/
+ module_1/
+ ...
+ optimizer_state/
+ optimizers_module_1/
+ ...
 
  Args:
  path: The path to the directory to save the state to.

rllib/core/learner/learner_group.py

@@ -475,10 +475,10 @@ def load_state(self, path: str) -> None:
  if not path.exists():
  raise ValueError(f"Path {path} does not exist.")
  path = str(path.absolute())
- assert len(self._workers) == self._worker_manager.num_healthy_actors()
  if self.is_local:
  self._learner.load_state(path)
  else:
+ assert len(self._workers) == self._worker_manager.num_healthy_actors()
  head_node_ip = socket.gethostbyname(socket.gethostname())
  workers = self._worker_manager.healthy_actor_ids()
 

rllib/core/rl_module/rl_module.py

@@ -186,7 +186,7 @@ def to_dict(self):
 
  """
  catalog_class_path = (
- serialize_type(type(self.catalog_class)) if self.catalog_class else ""
+ serialize_type(self.catalog_class) if self.catalog_class else ""
  )
  return {
  "observation_space": gym_space_to_dict(self.observation_space),

rllib/evaluation/rollout_worker.py

@@ -134,7 +134,7 @@ def _update_env_seed_if_necessary(
 
  NOTE: this may not work with remote environments (issue #18154).
  """
- if not seed:
+ if seed is None:
  return
 
  # A single RL job is unlikely to have more than 10K

rllib/tests/test_algorithm_save_load_checkpoint_learner.py

@@ -0,0 +1,128 @@
+import tempfile
+import unittest
+
+import ray
+from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
+from ray.rllib.utils.test_utils import check, framework_iterator
+from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.utils.metrics.learner_info import LEARNER_INFO, LEARNER_STATS_KEY
+
+
+algorithms_and_configs = {
+ "PPO": (PPOConfig().training(train_batch_size=2, sgd_minibatch_size=2))
+}
+
+
+@ray.remote
+def save_and_train(algo_cfg: AlgorithmConfig, env: str, tmpdir):
+ """Create an algo, checkpoint it, then train for 2 iterations.
+
+ Note: This function uses a seeded algorithm that can modify the global random state.
+ Running it multiple times in the same process can affect other algorithms.
+ Making it a Ray task runs it in a separate process and prevents it from
+ affecting other algorithms' random state.
+
+ Args:
+ algo_cfg: The algorithm config to build the algo from.
+ env: The gym genvironment to train on.
+ tmpdir: The temporary directory to save the checkpoint to.
+
+ Returns:
+ The learner stats after 2 iterations of training.
+ """
+ algo_cfg = (
+ algo_cfg.training(_enable_learner_api=True)
+ .rl_module(_enable_rl_module_api=True)
+ .rollouts(num_rollout_workers=0)
+ # setting min_time_s_per_iteration=0 and min_sample_timesteps_per_iteration=1
+ # to make sure that we get results as soon as sampling/training is done at
+ # least once
+ .reporting(min_time_s_per_iteration=0, min_sample_timesteps_per_iteration=1)
+ .debugging(seed=10)
+ )
+ algo = algo_cfg.environment(env).build()
+
+ tmpdir = str(tmpdir)
+ algo.save_checkpoint(tmpdir)
+ for _ in range(2):
+ results = algo.train()
+ return results["info"][LEARNER_INFO][DEFAULT_POLICY_ID][LEARNER_STATS_KEY]
+
+
+@ray.remote
+def load_and_train(algo_cfg: AlgorithmConfig, env: str, tmpdir):
+ """Loads the checkpoint saved by save_and_train and trains for 2 iterations.
+
+ Note: This function uses a seeded algorithm that can modify the global random state.
+ Running it multiple times in the same process can affect other algorithms.
+ Making it a Ray task runs it in a separate process and prevents it from
+ affecting other algorithms' random state.
+
+ Args:
+ algo_cfg: The algorithm config to build the algo from.
+ env: The gym genvironment to train on.
+ tmpdir: The temporary directory to save the checkpoint to.
+
+ Returns:
+ The learner stats after 2 iterations of training.
+
+ """
+ algo_cfg = (
+ algo_cfg.training(_enable_learner_api=True)
+ .rl_module(_enable_rl_module_api=True)
+ .rollouts(num_rollout_workers=0)
+ # setting min_time_s_per_iteration=0 and min_sample_timesteps_per_iteration=1
+ # to make sure that we get results as soon as sampling/training is done at
+ # least once
+ .reporting(min_time_s_per_iteration=0, min_sample_timesteps_per_iteration=1)
+ .debugging(seed=10)
+ )
+ algo = algo_cfg.environment(env).build()
+ tmpdir = str(tmpdir)
+ algo.load_checkpoint(tmpdir)
+ for _ in range(2):
+ results = algo.train()
+ return results["info"][LEARNER_INFO][DEFAULT_POLICY_ID][LEARNER_STATS_KEY]
+
+
+class TestAlgorithmWithLearnerSaveAndRestore(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls) -> None:
+ ray.init()
+
+ @classmethod
+ def tearDowClass(cls) -> None:
+ ray.shutdown()
+
+ def test_save_and_restore(self):
+ for algo_name in algorithms_and_configs:
+ config = algorithms_and_configs[algo_name]
+ for _ in framework_iterator(config, frameworks=["torch", "tf2"]):
+ with tempfile.TemporaryDirectory() as tmpdir:
+ # create an algorithm, checkpoint it, then train for 2 iterations
+ ray.get(save_and_train.remote(config, "CartPole-v1", tmpdir))
+ # load that checkpoint into a new algorithm and train for 2
+ # iterations
+ results_algo_2 = ray.get(
+ load_and_train.remote(config, "CartPole-v1", tmpdir)
+ )
+
+ # load that checkpoint into another new algorithm and train for 2
+ # iterations
+ results_algo_3 = ray.get(
+ load_and_train.remote(config, "CartPole-v1", tmpdir)
+ )
+
+ # check that the results are the same across loaded algorithms
+ # they won't be the same as the first algorithm since the random
+ # state that is used for each algorithm is not preserved across
+ # checkpoints.
+ check(results_algo_3, results_algo_2)
+
+
+if __name__ == "__main__":
+ import sys
+ import pytest
+
+ sys.exit(pytest.main(["-v", __file__]))

rllib/utils/checkpoints.py

@@ -29,7 +29,11 @@
 
 # 1.1: Same as 1.0, but has a new "format" field in the rllib_checkpoint.json file
 # indicating, whether the checkpoint is `cloudpickle` (default) or `msgpack`.
+
+# 1.2: Introduces the checkpoint for the new Learner API if the Learner api is enabled.
+
 CHECKPOINT_VERSION = version.Version("1.1")
+CHECKPOINT_VERSION_LEARNER = version.Version("1.2")
 
 
 @PublicAPI(stability="alpha")
@@ -102,15 +106,15 @@ def get_checkpoint_info(checkpoint: Union[str, Checkpoint]) -> Dict[str, Any]:
  rllib_checkpoint_info["checkpoint_version"]
  )
  info.update(rllib_checkpoint_info)
-
- # No rllib_checkpoint.json file present: Warn and continue trying to figure out
- # checkpoint info ourselves.
- if log_once("no_rllib_checkpoint_json_file"):
- logger.warning(
- "No `rllib_checkpoint.json` file found in checkpoint directory "
- f"{checkpoint}! Trying to extract checkpoint info from other files "
- f"found in that dir."
- )
+ else:
+  # No rllib_checkpoint.json file present: Warn and continue trying to figure
+  # out checkpoint info ourselves.
+  if log_once("no_rllib_checkpoint_json_file"):
+  logger.warning(
+  "No `rllib_checkpoint.json` file found in checkpoint directory "
+  f"{checkpoint}! Trying to extract checkpoint info from other files "
+  f"found in that dir."
+  )
 
  # Policy checkpoint file found.
  for extension in ["pkl", "msgpck"]:
@@ -222,7 +226,10 @@ def convert_to_msgpack_checkpoint(
  state["worker"]["is_policy_to_train"] = NOT_SERIALIZABLE
 
  # Add RLlib checkpoint version (as string).
- state["checkpoint_version"] = str(CHECKPOINT_VERSION)
+ if state["config"]["_enable_learner_api"]:
+ state["checkpoint_version"] = str(CHECKPOINT_VERSION_LEARNER)
+ else:
+ state["checkpoint_version"] = str(CHECKPOINT_VERSION)
 
  # Write state (w/o policies) to disk.
  state_file = os.path.join(msgpack_checkpoint_dir, "algorithm_state.msgpck")