新 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 默认启用新 API 栈以支持所有 RLlib 算法。

注意

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

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

请注意，配置旧 API 栈算法和其新栈对应版本之间存在一些其他差异。请仔细阅读以下各节，确保你已正确翻译了 respective 设置。删除新栈不支持或不需要的设置。

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 设置。要将 RLModule 放置在 Learner 端的一个或多个 GPU 上，请执行以下操作：

# 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 设置决定了你的 Algorithm 的 LearnerGroup 中有多少个远程 Learner 工作节点。如果将此参数设置为 0，则你的 LearnerGroup 只包含一个本地 Learner，该 Learner 在主进程上运行并共享其计算资源，通常是 1 个 CPU。因此，对于像 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，并且想将整个 NN 逻辑迁移到新栈，请参阅 ModelV2 到 RLModule 以获取迁移说明。

学习率和系数调度#

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

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

例如，要创建一个学习率调度，使其从 1e-5 开始，然后在 100 万个时间步后增加到 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 开始，然后在 100 万个时间步后增加到 0.1，然后在 100 万个时间步后突然降至 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 工作节点，因此不使用此方法。

它允许你指定：

通过 .learners(num_learners=...) 指定 Learner 工作节点的数量。
每个 Learner 的资源；对于 GPU 训练使用 .learners(num_gpus_per_learner=1)，对于 CPU 训练使用 .learners(num_gpus_per_learner=0)。
要使用的自定义 Learner 类。有关更多详细信息，请参阅此示例。
要为自定义 Learner 设置的配置字典：.learners(learner_config_dict={...})。请注意，每个 Learner 都可以通过 self.config 访问整个 AlgorithmConfig 对象，但设置 learner_config_dict 是一个方便的方法，可以避免创建全新的 AlgorithmConfig 子类，仅仅是为了支持自定义 Learner 类的一些额外设置。

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 的任何示例或调优示例脚本来实现这一点，只需设置命令行参数：--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 团队已将其重命名为 RLlibCallback。你可以继续采用此方法在新 API 栈上，并将自定义子类传递给配置，如下所示：

# config.callbacks(YourCallbacksClass)

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

以下是这些类型的 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 类，并希望将整个 NN 架构以及可能的动作分布逻辑转换为新 API 栈，请参阅 RL Modules 以及本节。

此外，请参阅以下示例脚本，了解如何编写自定义的包含 CNN 的 RL Module以及如何编写自定义的包含 LSTM 的 RL Module。

有多种选项可以将旧 API 栈中现有的自定义 ModelV2 迁移到新 API 栈的 RLModule：

将你的 ModelV2 代码移到一个新的、自定义的 RLModule 类中。有关详细信息，请参阅RL Modules。
使用你从旧 API 栈训练运行中获得的 Algorithm 检查点或 Policy 检查点，并使用此检查点与新栈 RL Module 便利包装器。
使用旧 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。

以下示例脚本展示了如何编写：

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

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

自定义连接器#

如果你正在使用旧 API 栈的自定义连接器，请将你的逻辑移到新的 ConnectorV2 API 中。将你的代理连接器翻译为从环境到模块的 ConnectorV2 片段，并将你的动作连接器翻译为从模块到环境的 ConnectorV2 片段。

ConnectorV2 文档正在开发中。

以下是一些关于如何为不同管道编写 ConnectorV2 片段的示例：