使用Ray RLlib构建自定义强化学习环境

强化学习是机器学习中一个令人兴奋的领域，它允许AI代理通过与环境交互来学习最优策略。本文将介绍如何使用Ray RLlib框架创建自定义强化学习环境并进行训练。无论您是强化学习新手还是寻求高性能分布式训练的专家，这篇指南都能帮您快速上手。

准备工作

在开始前，我们需要安装必要的依赖库：

pip install "ray[rllib]" gym torch

注意：如果使用macOS/Linux，可能还需要确保系统有C++编译工具链。若需要Atari/Mujoco等更复杂的环境，则需要安装额外的依赖。

自定义环境

首先，我们创建一个最小可运行的自定义环境MyCustomEnv，它完全兼容OpenAI Gym的接口规范（实现reset、step方法，并定义observation_space和action_space）：

import gym
from gym import spaces
import numpy as np

class MyCustomEnv(gym.Env):
    def __init__(self, config=None):
        super(MyCustomEnv, self).__init__()
        self.config = config or {}

        # 假设观测空间是一个长度为4的向量
        self.observation_space = spaces.Box(low=-1.0, high=1.0, shape=(4,), dtype=np.float32)
        # 假设动作空间离散，有2个动作可选
        self.action_space = spaces.Discrete(2)

        self.state = None

    def reset(self):
        # 重置环境状态并返回初始观测
        self.state = np.zeros(4, dtype=np.float32)
        return self.state

    def step(self, action):
        # 简单示例逻辑：action=1时+1奖励，否则0
        reward = 1.0 if action == 1 else 0.0
        # 用随机噪声更新一下 state，让它看起来不像纯0
        self.state = self.state + np.random.randn(4) * 0.01

        # 让episode在小概率下结束
        done = np.random.rand() < 0.05
        info = {}

        return self.state, reward, done, info

自定义神经网络模型

如果RLlib自带的默认网络结构不能满足需求，可以自定义PyTorch模型。只需继承TorchModelV2并实现必要的方法：

import torch
import torch.nn as nn
import torch.nn.functional as F

from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.utils.annotations import override

class MyCustomDNNModel(TorchModelV2, nn.Module):
    def __init__(self, obs_space, action_space, num_outputs, model_config, name):
        TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name)
        nn.Module.__init__(self)

        # 假设观测是shape=(4,)的向量
        input_size = obs_space.shape[0]
        hidden_size = 64  # 可以根据需求调整

        # 构建一个简单两层全连接网络
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)

        # 策略输出层和价值函数输出层
        self.policy_layer = nn.Linear(hidden_size, num_outputs)
        self.value_layer = nn.Linear(hidden_size, 1)

        self._value_out = None  # 用于存放价值函数输出

    @override(ModelV2)
    def forward(self, input_dict, state, seq_lens):
        # 处理观测
        obs = input_dict["obs"].float()
        x = F.relu(self.fc1(obs))
        x = F.relu(self.fc2(x))

        # 输出策略 (logits) 与价值 (self._value_out)
        logits = self.policy_layer(x)
        self._value_out = self.value_layer(x).squeeze(1)  # shape: [B]

        return logits, state

    @override(ModelV2)
    def value_function(self):
        # RLlib在计算价值函数损失时会自动调用这个方法
        return self._value_out

注册环境和模型

使用Ray提供的注册机制，将自定义的环境和模型注册到系统中：

from ray.tune.registry import register_env, register_custom_model

def my_env_creator(env_config):
    return MyCustomEnv(env_config)

register_env("MyEnv-v0", my_env_creator)
register_custom_model("my_dnn_model", MyCustomDNNModel)

注册后，可以在配置中通过名称引用它们，非常方便。

训练脚本编写

下面是一个完整的训练脚本示例，使用PPO算法对我们的自定义环境进行训练：

import ray
from ray.rllib.agents import ppo

if __name__ == "__main__":
    # 1. 启动Ray分布式框架
    ray.init()

    # 2. 配置训练参数
    config = {
        "env": "MyEnv-v0",           # 使用刚才注册的自定义环境
        "num_workers": 2,            # 并行worker数量；可根据CPU核心数灵活调整
        "framework": "torch",        # 使用PyTorch
        "model": {
            "custom_model": "my_dnn_model",  # 自定义的网络结构
        },
        # 训练超参数（可根据情况调整）
        "train_batch_size": 4000,
        "sgd_minibatch_size": 128,
        "lr": 1e-3,
    }

    # 3. 创建PPO的Trainer实例
    trainer = ppo.PPOTrainer(config=config)

    # 4. 多轮训练
    for i in range(10):
        result = trainer.train()
        print(f"轮次 {i}, 平均奖励: {result['episode_reward_mean']}")

    # 5. 清理资源
    ray.shutdown()

总结与进阶

通过以上步骤，我们完成了使用Ray RLlib构建自定义强化学习环境并训练模型的全过程。Ray的强大之处在于它自动处理了并行数据采集和分布式训练的复杂性，使我们能专注于环境设计和算法调优。

Ray RLlib的优势：

• 高效并行：自动创建多个Worker并行收集样本数据

• 灵活定制：支持自定义环境、网络模型和算法参数

• 分布式训练：可轻松扩展到多台机器进行训练

• 丰富的算法库：不仅支持PPO，还包括DQN、SAC、IMPALA等多种主流算法

希望这篇指南能帮助您快速上手Ray RLlib，开启强化学习的探索之旅！