新 API 栈迁移指南

注意

Ray 2.40 默认使用 RLlib 的新 API 栈。Ray 团队已基本完成将算法、示例脚本和文档迁移到新代码库的工作。

如果你仍在使用旧 API 栈，请参阅新 API 栈迁移指南了解如何迁移的详细信息。

本页将逐步解释如何将现有的旧 API 栈 RLlib 类和代码转换并迁移到 RLlib 的新 API 栈。

什么是新 API 栈？

新 API 栈是在彻底重写核心 RLlib API 的基础上形成的，将用户面临的关键类从十几个减少到仅少数几个，同时不丢失任何功能。在设计这些新接口时，Ray 团队严格遵循以下原则

• 类必须可在 RLlib 外部使用。

• 关注点分离。尝试回答：“何时、由谁应该完成什么？”并为每个类分配尽可能少、不重叠且定义明确的任务。

• 提供细粒度的模块化、完全的互操作性以及无障碍的类可插拔性。

• 尽可能使用广泛接受的第三方标准和 API。

应用上述原则，Ray 团队将普通 RLlib 用户需要了解的重要类从旧栈的八个减少到新栈的五个。核心新 API 栈类包括

• RLModule，它取代了 ModelV2 和 PolicyMap API

• Learner，它取代了 RolloutWorker 和部分 Policy

• SingleAgentEpisode 和 MultiAgentEpisode，它们取代了 ViewRequirement、SampleCollector、Episode 和 EpisodeV2

• ConnectorV2，它取代了 Connector 和部分 RolloutWorker 和 Policy

AlgorithmConfig 和 Algorithm API 保持不变。这些类在旧栈上已经是成熟的 API 了。

注意

尽管新 API 栈仍为 TensorFlow 提供初步支持，但 RLlib 只支持单一的深度学习框架 PyTorch，完全放弃了对 TensorFlow 的支持。但请注意，Ray 团队继续将 RLlib 设计为框架无关的，未来可能会增加对其他框架的支持。

检查你的 AlgorithmConfig

对于所有 RLlib 算法，RLlib 默认启用新 API 栈。

注意

要停用新 API 栈并切换回旧 API 栈，请在你的 AlgorithmConfig 对象中使用 api_stack() 方法，如下所示

config.api_stack(
    enable_rl_module_and_learner=False,
    enable_env_runner_and_connector_v2=False,
)

请注意，配置旧 API 栈算法与新栈算法之间还有一些其他差异。请阅读以下章节，确保你正在转换相应的设置。删除新栈不支持或不需要的设置。

AlgorithmConfig.framework()

尽管新 API 栈仍为 TensorFlow 提供初步支持，但 RLlib 只支持单一的深度学习框架 PyTorch。

新 API 栈弃用了以下与框架相关的设置

# Make sure you always set the framework to "torch"...
config.framework("torch")

# ... and drop all tf-specific settings.
config.framework(
    eager_tracing=True,
    eager_max_retraces=20,
    tf_session_args={},
    local_tf_session_args={},
)

AlgorithmConfig.resources()

Ray 团队弃用了 num_gpus 和 _fake_gpus 设置。要在 Learner 端的一个或多个 GPU 上放置你的 RLModule，请执行以下操作

# The following setting is equivalent to the old stack's `config.resources(num_gpus=2)`.
config.learners(
    num_learners=2,
    num_gpus_per_learner=1,
)

提示

num_learners 设置决定了你的算法的 LearnerGroup 中有多少个远程 Learner 工作进程。如果你将此参数设置为 0，则你的 LearnerGroup 只包含一个在主进程上运行并共享计算资源（通常是 1 个 CPU）的本地 Learner。因此，对于像 IMPALA 或 APPO 这样的异步算法，此设置应始终大于 0。

此处查看如何使用分数 GPU 进行训练的示例。另请注意，对于分数 GPU，你应始终将 num_learners 设置为 0 或 1。

如果 GPU 不可用，但你想以多 CPU 方式使用多个 Learner 进行学习，你可以执行以下操作

config.learners(
    num_learners=2,  # or >2
    num_cpus_per_learner=1,  # <- default
    num_gpus_per_learner=0,  # <- default
)

Ray 团队将设置 num_cpus_for_local_worker 重命名为 num_cpus_for_main_process。

config.resources(num_cpus_for_main_process=0)  # default is 1

AlgorithmConfig.training()

训练批大小

由于新 API 栈的 Learner 工作进程架构，训练可能通过 n 个 Learner 工作进程以分布式方式进行，因此 RLlib 提供每个单独 Learner 的训练批大小。不再使用 train_batch_size 设置

config.training(
    train_batch_size_per_learner=512,
)

即使通过 config.learners(num_learners=...) 增加 Learner 的数量，你也不需要更改此设置。

请注意，在 learner 维度进行扩展的一个经验法则是，随着 Learner 数量的增长，保持 train_batch_size_per_learner 值不变，并按如下方式增加学习率

lr = [original_lr] * ([num_learners] ** 0.5)

神经网络配置

旧栈的 config.training(model=...) 在新 API 栈中不再支持。相反，使用新的 rl_module() 方法来配置 RLlib 的默认 RLModule 或指定和配置自定义 RLModule。

请参阅RLModules API，这是一份通用指南，也解释了 config.rl_module() 方法的使用。

如果你有一个旧栈的 ModelV2 并想将整个神经网络逻辑迁移到新栈，请参阅 ModelV2 到 RLModule 获取迁移说明。

学习率和系数调度

如果你使用学习率或其他系数的调度（例如 PPO 中的 entropy_coeff 设置），请直接在相应的设置中提供调度信息。调度行为不再需要特定的独立设置。

定义调度时，提供一个 2 元组列表，其中第一个项是全局时间步长（在报告的指标中为 num_env_steps_sampled_lifetime），第二个项是学习率在该时间步长应达到的值。第一个 2 元组始终以时间步长 0 开始。请注意，RLlib 会在两个提供的时间步长之间进行线性插值。

例如，要创建一个学习率调度，使其以 1e-5 的值开始，然后在一百万个时间步长内增加到 1e-4，之后保持恒定，请执行以下操作

config.training(
    lr=[
        [0, 1e-5],  # <- initial value at timestep 0
        [1000000, 1e-4],  # <- final value at 1M timesteps
    ],
)

在前面的示例中，经过 50 万个时间步长后的值通过线性插值约为 5e-5。

另一个示例是创建一个熵系数调度，使其以 0.05 的值开始，然后在一百万个时间步长内增加到 0.1，并在达到一百万个时间步长后突然下降到 0，请执行以下操作

config.training(
    entropy_coeff=[
        [0, 0.05],  # <- initial value at timestep 0
        [1000000, 0.1],  # <- value at 1M timesteps
        [1000001, 0.0],  # <- sudden drop to 0.0 right after 1M timesteps
    ]
)

如果需要配置更复杂的学习率调度行为或将不同的调度器链式连接到流水线中，可以使用实验性的 _torch_lr_schedule_classes 配置属性。请参阅 此示例脚本 获取更多详细信息。请注意，此示例仅涵盖学习率调度，不包括任何其他系数。

AlgorithmConfig.learners()

此方法在旧 API 栈中不使用，因为旧栈不使用 Learner 工作进程。

它允许你指定

1. 通过 .learners(num_learners=...) 设置 Learner 工作进程的数量。

2. 每个 learner 的资源；对于 GPU 训练，使用 .learners(num_gpus_per_learner=1)；对于 CPU 训练，使用 .learners(num_gpus_per_learner=0)。

3. 你想使用的自定义 Learner 类。请参阅此示例了解更多详细信息。

4. 你想要为自定义 learner 设置的配置字典：.learners(learner_config_dict={...})。请注意，每个 Learner 都可以通过 self.config 访问完整的 AlgorithmConfig 对象，但设置 learner_config_dict 是一种便捷方式，可以避免仅为了支持自定义 Learner 类的一些额外设置而不得不创建一个全新的 AlgorithmConfig 子类。

AlgorithmConfig.env_runners()

# RolloutWorkers have been replace by EnvRunners. EnvRunners are more efficient and offer
# a more separation-of-concerns design and cleaner code.
config.env_runners(
    num_env_runners=2,  # use this instead of `num_workers`
)

# The following `env_runners` settings are deprecated and should no longer be explicitly
# set on the new stack:
config.env_runners(
    create_env_on_local_worker=False,
    sample_collector=None,
    enable_connectors=True,
    remote_worker_envs=False,
    remote_env_batch_wait_ms=0,
    preprocessor_pref="deepmind",
    enable_tf1_exec_eagerly=False,
    sampler_perf_stats_ema_coef=None,
)

提示

如果你想在 IDE 中调试 EnvRunners 中发生了什么，请设置 num_env_runners=0 并确保你在本地运行实验而不是通过 Ray Tune。要使用 RLlib 的任何 示例 或 tuned_example 脚本执行此操作，只需设置命令行参数：--no-tune --num-env-runners=0。

如果你之前使用了 observation_filter 设置，请执行以下转换

# For `observation_filter="NoFilter"`, don't set anything in particular. This is the default.

# For `observation_filter="MeanStdFilter"`, do the following:
from ray.rllib.connectors.env_to_module import MeanStdFilter

config.env_runners(
    env_to_module_connector=lambda env: MeanStdFilter(multi_agent=False),  # <- or True
)

提示

在样本收集期间是否进行探索的主要开关已移至 env_runners() 方法。请参阅 此处了解更多详细信息。

AlgorithmConfig.exploration()

在样本收集期间是否进行探索的主要开关已从已弃用的 AlgorithmConfig.exploration() 方法移至 env_runners()

它决定了你的 RLModule 在 EnvRunner 内部调用的方法是 _forward_exploration()（在 explore=True 的情况下）还是 _forward_inference()（在 explore=False 的情况下）。

config.env_runners(explore=True)  # <- or False

Ray 团队已弃用 exploration_config 设置。相反，在你的 RLModule 的重写 _forward_exploration() 方法中定义确切的探索行为，例如从分布中采样一个动作。

自定义回调

如果你在旧 API 栈中使用了自定义回调，你继承的是 DefaultCallbacks 类，Ray 团队将其重命名为 :py:class`~ray.rllib.callbacks.callbacks.RLlibCallback`。你可以在新 API 栈中继续使用这种方法，并将你的自定义子类传递给配置，如下所示

# config.callbacks(YourCallbacksClass)

但是，如果你重写了在 EnvRunner 端触发的方法，例如 on_episode_start/stop/step/等等...，你可能需要转换一些调用参数。

以下是此类 :py:class`~ray.rllib.callbacks.callbacks.RLlibCallback` 方法的一对一转换指南

from ray.rllib.callbacks.callbacks import RLlibCallback

class YourCallbacksClass(RLlibCallback):

    def on_episode_start(
        self,
        *,
        episode,
        env_runner,
        metrics_logger,
        env,
        env_index,
        rl_module,

        # Old API stack args; don't use or access these inside your method code.
        worker=None,
        base_env=None,
        policies=None,
        **kwargs,
    ):
        # The `SingleAgentEpisode` or `MultiAgentEpisode` that RLlib has just started.
        # See https://docs.rayai.org.cn/en/latest/rllib/single-agent-episode.html for more details:
        print(episode)

        # The `EnvRunner` class that collects the episode in question.
        # This class used to be a `RolloutWorker`. On the new stack, this class is either a
        # `SingleAgentEnvRunner` or a `MultiAgentEnvRunner` holding the gymnasium Env,
        # the RLModule, and the 2 connector pipelines, env-to-module and module-to-env.
        print(env_runner)

        # The MetricsLogger object on the EnvRunner (documentation is a WIP).
        print(metrics_logger.peek("episode_return_mean", default=0.0))

        # The gymnasium env that sample collection uses. Note that this env may be a
        # gymnasium.vector.VectorEnv.
        print(env)

        # The env index, in case of a vector env, that handles the `episode`.
        print(env_index)

        # The RL Module that this EnvRunner uses. Note that this module may be a "plain", single-agent
        # `RLModule`, or a `MultiRLModule` in the multi-agent case.
        print(rl_module)

# Change similarly:
# on_episode_created()
# on_episode_step()
# on_episode_end()

以下回调方法在新 API 栈中不再可用

• on_sub_environment_created()：新 API 栈使用 Farama 的 gymnasium 向量环境，RLlib 无法控制在每个单独环境索引创建时调用回调。

• on_create_policy()：此方法在新 API 栈中不再可用，因为它只由 RolloutWorker 调用。

• on_postprocess_trajectory()：新 API 栈不再触发和调用此方法，因为 ConnectorV2 流水线完全处理轨迹。有关 ConnectorV2 的文档正在开发中。

ModelV2 到 RLModule

如果你使用自定义 ModelV2 类并想将整个神经网络架构以及可能的动作分布逻辑迁移到新 API 栈，除了本节外，请参阅 RL Modules。

另请参阅这些示例脚本，了解如何编写包含自定义 CNN 的 RL Module 和如何编写包含自定义 LSTM 的 RL Module。

将现有自定义 ModelV2 从旧 API 栈转换为新 API 栈的 RLModule 有多种选项

1. 将你的 ModelV2 代码迁移到新的自定义 RLModule 类。请参阅 RL Modules 了解详细信息）。

2. 使用你在旧 API 栈训练运行中获得的 Algorithm 检查点或 Policy 检查点，并结合新栈 RL Module 便利包装器使用此检查点。

3. 使用你在旧 API 栈训练运行中获得的现有 AlgorithmConfig 对象，并结合新栈 RL Module 便利包装器使用。

在更复杂的场景中，你可能实现了自定义策略，以便修改模型和分布的构建行为。

转换 Policy.compute_actions_from_input_dict

这个旧 API 栈方法，以及 compute_actions 和 compute_single_action，直接转换为 _forward_inference() 和 _forward_exploration()。RLModule 指南解释了如何实现此方法。

转换 Policy.action_distribution_fn

要转换 action_distribution_fn，请编写以下自定义 RLModule 代码

相同的动作分布类

from ray.rllib.models.torch.torch_distributions import YOUR_DIST_CLASS


class MyRLModule(TorchRLModule):
    def setup(self):
        ...
        # Set the following attribute at the end of your custom `setup()`.
        self.action_dist_cls = YOUR_DIST_CLASS

不同的动作分布类

from ray.rllib.models.torch.torch_distributions import (
    YOUR_INFERENCE_DIST_CLASS,
    YOUR_EXPLORATION_DIST_CLASS,
    YOUR_TRAIN_DIST_CLASS,
)

    def get_inference_action_dist_cls(self):
        return YOUR_INFERENCE_DIST_CLASS

    def get_exploration_action_dist_cls(self):
        return YOUR_EXPLORATION_DIST_CLASS

    def get_train_action_dist_cls(self):
        return YOUR_TRAIN_DIST_CLASS

转换 Policy.action_sampler_fn

要转换 action_sampler_fn，请编写以下自定义 RLModule 代码

from ray.rllib.models.torch.torch_distributions import YOUR_DIST_CLASS


class MyRLModule(TorchRLModule):

    def _forward_exploration(self, batch):
        computation_results = ...
        my_dist = YOUR_DIST_CLASS(computation_results)
        actions = my_dist.sample()
        return {Columns.ACTIONS: actions}

    # Maybe for inference, you would like to sample from the deterministic version
    # of your distribution:
    def _forward_inference(self, batch):
        computation_results = ...
        my_dist = YOUR_DIST_CLASS(computation_results)
        greedy_actions = my_dist.to_deterministic().sample()
        return {Columns.ACTIONS: greedy_actions}

Policy.compute_log_likelihoods

实现你的自定义 RLModule 的 _forward_train() 方法，并返回 Columns.ACTION_LOGP 键以及相应的动作对数概率，以便将此信息传递给你的损失函数，你的代码在调用 forward_train() 后会调用这些损失函数。然后，损失逻辑可以访问 Columns.ACTION_LOGP。

自定义损失函数和策略

如果你使用一个或多个自定义损失函数或自定义 (PyTorch) 优化器来训练模型，则需要将逻辑从旧栈的 Policy 类迁移到新 API 栈的 Learner 类中。

请参阅 Learner 了解如何编写自定义 Learner 的详细信息。

以下示例脚本展示了如何编写： - 一个简单的自定义损失函数 - 一个带有 2 个优化器且每个优化器学习率不同的自定义 Learner。

请注意，新 API 栈不支持 Policy 类。在旧栈中，这个类包含一个神经网络（在新 API 栈中是 RLModule）、一个旧栈连接器（在新 API 栈中是 ConnectorV2）以及一个或多个优化器和损失函数（在新 API 栈中是 Learner 类）。

RL Module API 比旧栈的 Policy API 灵活得多，并提供更清晰的关注点分离体验。与动作推理相关的内容在 EnvRunners 上运行，而与更新相关的内容在 Learner 工作进程上运行。它还提供了出色的可扩展性，允许在任何 Ray 集群中进行多 GPU 设置下的训练，并在 Anyscale 平台上进行多节点多 GPU 训练。

自定义连接器

如果你使用旧 API 栈的自定义连接器，请将你的逻辑迁移到新的 ConnectorV2 API。将你的代理连接器转换为 env-to-module ConnectorV2 片段，将你的动作连接器转换为 module-to-env ConnectorV2 片段。

有关 ConnectorV2 的文档正在开发中。

以下是一些关于如何为不同流水线编写 ConnectorV2 片段的示例

1. 观测帧堆叠.

2. 将最近的动作和奖励添加到 RL Module 的输入中.

3. 对所有观测进行均值-标准差过滤.

4. 将任何复杂的观测空间展平到一维空间.