Checkpointing

RLlib 为其所有主要类提供了一个强大的 checkpointing 系统，允许您将 Algorithm 实例及其子组件的状态保存到本地磁盘或云存储，并恢复先前运行的实验状态和单个子组件。该系统允许您从先前的状态继续训练模型，或将基本的 PyTorch 模型部署到生产环境中。

保存到和从磁盘或云存储恢复：使用 save_to_path() 方法将任何 Checkpointable() 组件或您的整个 Algorithm 的当前状态写入磁盘或云存储。要将已保存的状态加载回正在运行的组件或您的 Algorithm，请使用 restore_from_path() 方法。

Checkpoint 是磁盘上的一个目录，或者是一个 PyArrow 支持的云位置，例如 gcs 或 S3。它包含架构信息，例如用于创建新实例的类和构造函数参数，一个包含状态信息的 pickle 或 msgpack 文件，以及一个包含 Ray 版本、git 提交和 checkpoint 版本信息的、人类可读的 metadata.json 文件。

您可以使用 from_checkpoint() 方法从现有 checkpoint 生成一个新的 Algorithm 实例或其他子组件，例如 RLModule。例如，您可以将先前训练过的 RLModule（不包含任何其他 RLlib 组件）部署到生产环境中。

直接从 checkpoint 创建新实例：使用 classmethod from_checkpoint() 直接从 checkpoint 实例化对象。RLlib 首先使用保存的元数据创建一个原始 checkpointed 对象的精简实例，然后从 checkpoint 目录中的状态信息恢复其状态。

另一种可能性是只将某个子组件的状态加载到包含它的更高层级对象中。例如，您可能只想加载 Algorithm 中的 RLModule 的状态，而保持其他所有组件不变。

Checkpointable API

RLlib 通过 Checkpointable API 管理 checkpointing，该 API 提供了以下三个主要方法：

• save_to_path() 用于创建新的 checkpoint

• restore_from_path() 用于将状态从 checkpoint 加载到正在运行的对象中

• from_checkpoint() 用于从 checkpoint 创建新对象

到目前为止支持 Checkpointable API 的 RLlib 类有：

• 算法

• RLModule (以及 MultiRLModule)

• EnvRunner (因此，也包括 SingleAgentEnvRunner 和 MultiAgentEnvRunner)

• ConnectorV2 (因此，也包括 ConnectorPipelineV2)

• LearnerGroup

• Learner

使用 save_to_path() 创建新 checkpoint

您可以通过 save_to_path() 方法从已实例化的 RLlib 对象创建一个新的 checkpoint。

以下是两个示例，单智能体和多智能体，使用 Algorithm 类，展示了如何创建 checkpoints。

单智能体设置

from ray.rllib.algorithms.ppo import PPOConfig

# Configure and build an initial algorithm.
config = (
    PPOConfig()
    .environment("Pendulum-v1")
)
ppo = config.build()

# Train for one iteration, then save to a checkpoint.
print(ppo.train())
checkpoint_dir = ppo.save_to_path()
print(f"saved algo to {checkpoint_dir}")

多智能体设置

from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentPendulum
from ray.tune import register_env

register_env("multi-pendulum", lambda cfg: MultiAgentPendulum({"num_agents": 2}))

# Configure and build an initial algorithm.
multi_agent_config = (
    PPOConfig()
    .environment("multi-pendulum")
    .multi_agent(
        policies={"p0", "p1"},
        # Agent IDs are 0 and 1 -> map to p0 and p1, respectively.
        policy_mapping_fn=lambda aid, eps, **kw: f"p{aid}"
    )
)
ppo = multi_agent_config.build()

# Train for one iteration, then save to a checkpoint.
print(ppo.train())
multi_agent_checkpoint_dir = ppo.save_to_path()
print(f"saved multi-agent algo to {multi_agent_checkpoint_dir}")

注意

当您使用 Ray Tune 运行您的实验时，Tune 会在训练迭代次数匹配通过 Tune 配置的 checkpoint 频率时，自动在 Algorithm 实例上调用 save_to_path() 方法。Tune 创建这些 checkpoints 的默认位置是 ~/ray_results/[your experiment name]/[Tune trial name]/checkpoint_[sequence number]。

Checkpoint 版本

RLlib 使用 checkpoint 版本控制系统来确定如何从给定目录恢复 Algorithm 或任何子组件。

从 Ray 2.40 开始，您可以在所有 checkpoint 目录内的、人类可读的 metadata.json 文件中找到 checkpoint 版本。

同样从 Ray 2.40 开始，RLlib checkpoints 是向后兼容的。这意味着一个用 Ray 2.x 创建的 checkpoint 可以被 Ray 2.x+n 读取和处理，只要 x >= 40。Ray 团队通过 关于使用先前 Ray 版本创建的 checkpoints 的全面的 CI 测试 来确保向后兼容性。

Checkpoint 目录结构

在将 PPO 的状态保存在 checkpoint_dir 目录中，或者如果您使用 Ray Tune，则保存在 ~/ray_results/ 中的某个位置之后，目录会如下所示：

$ cd [your algo checkpoint dir]
$ ls -la
    .
    ..
    env_runner/
    learner_group/
    algorithm_state.pkl
    class_and_ctor_args.pkl
    metadata.json

Checkpoint 目录内的子目录，例如 env_runner/，暗示了子组件自己的 checkpoint 数据。例如，一个 Algorithm 总是会保存它的 EnvRunner 状态和 LearnerGroup 状态。

注意

每个子组件的目录本身包含一个 metadata.json 文件、一个 class_and_ctor_args.pkl 文件，以及一个 pickle 或 msgpack 状态文件，这些文件与其在主 algorithm checkpoint 目录中的对应项具有相同的功能。例如，在 learner_group/ 子目录内，您会找到 LearnerGroup 自己的架构、状态和元信息。

$ cd env_runner/
$ ls -la
.
..
state.pkl
class_and_ctor_args.pkl
metadata.json

有关详细信息，请参阅 RLlib component tree。

metadata.json 文件仅供您方便之用，RLlib 并不需要它。

注意

metadata.json 文件包含有关用于创建 checkpoint 的 Ray 版本、Ray commit、RLlib checkpoint 版本以及同一目录中的状态信息文件和构造函数信息文件名称的信息。

$ more metadata.json
{
    "class_and_ctor_args_file": "class_and_ctor_args.pkl",
    "state_file": "state",
    "ray_version": ..,
    "ray_commit": ..,
    "checkpoint_version": "2.1"
}

class_and_ctor_args.pkl 文件存储了创建“全新”对象（不包含任何特定状态）所需的元信息。正如文件名所示，此信息包含已保存对象的类及其构造函数参数和关键字参数。RLlib 在调用 from_checkpoint() 时使用此文件来创建初始新对象。

最后，.._state.[pkl|msgpack] 文件包含已保存对象的 pickled 或 msgpacked 状态字典。RLlib 在保存 checkpoint 时，通过调用对象的 get_state() 方法来获取此状态字典。

注意

对 msgpack 格式的 checkpoints 的支持是实验性的，但将来可能会成为默认设置。与 pickle 不同，msgpack 的优点是独立于 python 版本，因此允许用户从使用旧 python 版本生成的旧 checkpoints 中恢复实验和模型状态。

Ray 团队正在努力将状态与架构完全分离，这意味着所有状态信息都应进入 state.msgpack 文件（与 python 版本无关），而所有架构信息都应进入 class_and_ctor_args.pkl 文件（仍然依赖于 python 版本）。在加载 checkpoint 时，用户需要提供后者的架构部分。

在此处查看一个更详细地说明此点的示例.

RLlib 组件树

以下是 RLlib 组件树的结构，显示了在更高的 checkpoint 内，您可以通过哪个名称访问子组件自己的 checkpoint。最高级别是 Algorithm 类。

algorithm/
    learner_group/
        learner/
            rl_module/
                default_policy/  # <- single-agent case
                [module ID 1]/  # <- multi-agent case
                [module ID 2]/  # ...
    env_runner/
        env_to_module_connector/
        module_to_env_connector/

注意

env_runner/ 子组件目前不包含 RLModule checkpoint 的副本，因为它已在 learner/ 下保存。Ray 团队正在努力解决此问题，可能通过软链接来避免重复文件和不必要的磁盘使用。

使用 from_checkpoint 从 checkpoint 创建实例

一旦您拥有一个经过训练的 Algorithm 或其任何 子组件 的 checkpoint，您就可以直接从该 checkpoint 重新创建新对象。

以下是两个示例：

从 checkpoint 创建新的 Algorithm

要从 checkpoint 重新创建整个 Algorithm 实例，您可以执行以下操作：

# Import the correct class to create from scratch using the checkpoint.
from ray.rllib.algorithms.algorithm import Algorithm

# Use the already existing checkpoint in `checkpoint_dir`.
new_ppo = Algorithm.from_checkpoint(checkpoint_dir)
# Confirm the `new_ppo` matches the originally checkpointed one.
assert new_ppo.config.env == "Pendulum-v1"

# Continue training.
new_ppo.train()

从 Algorithm checkpoint 创建新的 RLModule

从 Algorithm checkpoint 创建新的 RLModule 对于将训练好的模型部署到生产环境或在训练过程中在单独的进程中对其进行评估非常有用。要仅从 algorithm 的 checkpoint 重新创建 RLModule，您可以执行以下操作：

from pathlib import Path
import torch

# Import the correct class to create from scratch using the checkpoint.
from ray.rllib.core.rl_module.rl_module import RLModule

# Use the already existing checkpoint in `checkpoint_dir`, but go further down
# into its subdirectory for the single RLModule.
# See the preceding section on "RLlib component tree" for the various elements in the RLlib
# component tree.
rl_module_checkpoint_dir = Path(checkpoint_dir) / "learner_group" / "learner" / "rl_module" / "default_policy"

# Now that you have the correct subdirectory, create the actual RLModule.
rl_module = RLModule.from_checkpoint(rl_module_checkpoint_dir)

# Run a forward pass to compute action logits.
# Use a dummy Pendulum observation tensor (3d) and add a batch dim (B=1).
results = rl_module.forward_inference(
    {"obs": torch.tensor([0.5, 0.25, -0.3]).unsqueeze(0).float()}
)
print(results)

请参阅此 关于训练后运行策略推断的示例 和此 关于使用 LSTM 运行策略推断的示例。

提示

由于您的 RLModule 也是一个 PyTorch Module，您可以轻松地将模型导出到 ONNX、IREE 或其他易于部署的格式。有关更多详细信息，请参阅此 支持 ONNX 的示例脚本。

使用 restore_from_path 从 checkpoint 恢复状态

通常，save_to_path() 和 from_checkpoint() 方法足以创建 checkpoints 并从中重新创建实例。

但是，有时您已经拥有一个正在运行的已实例化对象，并且希望将另一个状态“加载”到其中。例如，考虑通过多智能体训练来训练两个 RLModule 网络，并以自对弈的方式相互对弈。一段时间后，您希望在不中断实验的情况下，将其中一个 RLModules 替换为一个您在很久以前已保存到磁盘或云存储的第三个。

这时 restore_from_path() 方法就派上用场了。它将状态加载到已运行的对象中，例如您的 Algorithm，或该对象的一个子组件中，例如您 Algorithm 中的特定 RLModule。

继续训练

直接使用 RLlib 时（即不使用 Ray Tune），将状态加载到正在运行的实例中的问题很简单。

# Recreate the preceding PPO from the config.
new_ppo = config.build()

# Load the state stored previously in `checkpoint_dir` into the
# running algorithm instance.
new_ppo.restore_from_path(checkpoint_dir)

# Run another training iteration.
new_ppo.train()

使用 Ray Tune 继续训练

然而，通过 Ray Tune 运行时，您无法直接访问 Algorithm 对象或其任何子组件。您可以使用 RLlib 的 callbacks API 来注入自定义代码并解决此问题。

另请参阅此 关于如何使用不同配置继续训练的示例。

from ray import tune

# Reuse the preceding PPOConfig (`config`).
# Inject custom callback code that runs right after algorithm's initialization.
config.callbacks(
    on_algorithm_init=(
        lambda algorithm, _dir=checkpoint_dir, **kw: algorithm.restore_from_path(_dir)
    ),
)

# Run the experiment, continuing from the checkpoint, through Ray Tune.
results = tune.Tuner(
    config.algo_class,
    param_space=config,
    run_config=tune.RunConfig(stop={"num_env_steps_sampled_lifetime": 8000})
).fit()

替换一个 RLModule 并继续多智能体训练

在 关于 save_to_path 的前一节 中，您使用 default_policy ModuleID 创建了一个单智能体 checkpoint，并使用两个 ModuleIDs，p0 和 p1，创建了一个多智能体 checkpoint。

以下是如何继续训练多智能体实验，但将 p1 替换为单智能体实验中 default_policy 的状态。您可以使用 RLlib 的 callbacks API 将自定义代码注入 Ray Tune 实验。

# Reuse the preceding multi-agent PPOConfig (`multi_agent_config`).

# But swap out ``p1`` with the state of the ``default_policy`` from the
# single-agent run, using a callback and the correct path through the
# RLlib component tree:
multi_rl_module_component_tree = "learner_group/learner/rl_module"

# Inject custom callback code that runs right after algorithm's initialization.

def _on_algo_init(algorithm, **kwargs):
    algorithm.restore_from_path(
        # Checkpoint was single-agent (has "default_policy" subdir).
        path=Path(checkpoint_dir) / multi_rl_module_component_tree / "default_policy",
        # Algo is multi-agent (has "p0" and "p1" subdirs).
        component=multi_rl_module_component_tree + "/p1",
    )

# Inject callback.
multi_agent_config.callbacks(on_algorithm_init=_on_algo_init)

# Run the experiment through Ray Tune.
results = tune.Tuner(
    multi_agent_config.algo_class,
    param_space=multi_agent_config,
    run_config=tune.RunConfig(stop={"num_env_steps_sampled_lifetime": 8000})
).fit()