ray.rllib-入门实践-12：自定义policy

        在本博客开始之前，先厘清一下几个概念之间的区别与联系：env,  agent,  model, algorithm， policy. 

        强化学习由两部分组成： 环境(env)和智能体(agent)。环境（env）提供观测值和奖励; agent读取观测值，输出动作或决策。agent是algorithm的类对象。 policy是algorithm的子类， 比如ppo, dqn等。因此，自定义policy本质上是自定义algorithm.  algorithm 主要由两部分组成： 网络结构（model）和损失函数(loss)。 网络结构（model）的自定义由上一个博客ray.rllib-入门实践-11： 自定义模型/网络 进行了介绍：在alrorithm外创建新的model类， 通过 AlgorithmConfig类传入algorithm。因此， c从实际操作上， 自定义algorithm就变成了自定义algorithm 的 loss.

        因此，本博客所提到的自定义policy， 本质上就是继承一个Algorithm， 并修改它的loss函数。

        与之前介绍的自定义env, 自定义model一样， 自定义policy也包含三个步骤：

        1. 继承某个Policy, 创建一个新Policy类， 修改它的损失函数。

        2. 把自己的Policy封装为一个Algorithm, 使ray可识别

        3. 配置使用自己的Policy.

环境配置：

        torch==2.5.1

        ray==2.10.0

        ray[rllib]==2.10.0

        ray[tune]==2.10.0

        ray[serve]==2.10.0

        numpy==1.23.0

        python==3.9.18

一、 自定义 policy

import torch 
import gymnasium as gym 
from gymnasium import spaces
from ray.rllib.utils.annotations import override
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.algorithms.ppo import PPO, PPOConfig, PPOTorchPolicy
from typing import Dict, List, Type, Union
from ray.rllib.models.action_dist import ActionDistribution
from ray.rllib.policy.sample_batch import SampleBatch

## 1. 自定义 policy， 主要是改变 policy 的 loss 的计算  # 神经网络的损失函数
class MY_PPOTorchPolicy(PPOTorchPolicy):
    """PyTorch policy class used with PPO."""
    def __init__(self, observation_space:gym.spaces.Box, action_space:gym.spaces.Box, config:PPOConfig): 
        PPOTorchPolicy.__init__(self,observation_space,action_space,config)
        ## PPOTorchPolicy 内部对 PPOConfig 格式的config 执行了to_dict()操作，后面可以以 dict 的形式使用 config

    @override(PPOTorchPolicy) 
    def loss(self,model: ModelV2,dist_class: Type[ActionDistribution],train_batch: SampleBatch):
        ## 原始损失
        original_loss = super().loss(model, dist_class, train_batch) # PPO原来的损失函数, 也可以完全自定义新的loss函数， 但是非常不建议。

        ## 新增自定义损失，这里以正则化损失作为示例
        addiontial_loss = torch.tensor(0.0) ## 自己定义的loss
        addiontial_loss = torch.tensor(0.)
        for param in model.parameters():
            addiontial_loss += torch.norm(param)
        ## 得到更新后的损失
        new_loss = original_loss + 0.01 * addiontial_loss
        return new_loss

二、 把自己的policy封装在一个算法中

## 2. 把自己的 policy 封装在算法中： 
##    继承自PPO， 创建一个新的算法类， 默认调用的是自定义的policy
class MY_PPO(PPO):
    ## 重写 get_default_policy_class 函数， 使其返回自定义的policy 
    def get_default_policy_class(self, config):
        return MY_PPOTorchPolicy

三、使用自己的策略创建智能体，执行训练

配置方法1：

## 三、使用自己的策略创建智能体，执行训练
## method-1
from ray.tune.logger import pretty_print 
config = PPOConfig(algo_class = MY_PPO) ## 配置使用自己的算法
config = config.environment("CartPole-v1")
config = config.rollouts(num_rollout_workers=2)
config = config.framework(framework="torch") 
algo = config.build()
result = algo.train()
print(pretty_print(result))

配置方法2：

## 3. 使用新策略执行训练
## method-2
from ray.tune.logger import pretty_print 
config = PPOConfig()
config = config.environment("CartPole-v1")
config = config.rollouts(num_rollout_workers=2)
config = config.framework(framework="torch") 
algo = MY_PPO(config=config,)  ## 在这里使用自己的policy
result = algo.train()
print(pretty_print(result))

四、代码汇总：

import torch 
import gymnasium as gym 
from gymnasium import spaces
from ray.rllib.utils.annotations import override
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.algorithms.ppo import PPO, PPOConfig, PPOTorchPolicy
from typing import Dict, List, Type, Union
from ray.rllib.models.action_dist import ActionDistribution
from ray.rllib.policy.sample_batch import SampleBatch
from ray.tune.logger import pretty_print 

## 1. 自定义 policy， 主要是改变 policy 的 loss 的计算  # 神经网络的损失函数
class MY_PPOTorchPolicy(PPOTorchPolicy):
    """PyTorch policy class used with PPO."""
    def __init__(self, observation_space:gym.spaces.Box, action_space:gym.spaces.Box, config:PPOConfig): 
        PPOTorchPolicy.__init__(self,observation_space,action_space,config)
        ## PPOTorchPolicy 内部对 PPOConfig 格式的config 执行了to_dict()操作，后面可以以 dict 的形式使用 config

    @override(PPOTorchPolicy) 
    def loss(self,model: ModelV2,dist_class: Type[ActionDistribution],train_batch: SampleBatch):
        ## 原始损失
        original_loss = super().loss(model, dist_class, train_batch) # PPO原来的损失函数, 也可以完全自定义新的loss函数， 但是非常不建议。

        ## 新增自定义损失，这里以正则化损失作为示例
        addiontial_loss = torch.tensor(0.0) ## 自己定义的loss
        addiontial_loss = torch.tensor(0.)
        for param in model.parameters():
            addiontial_loss += torch.norm(param)
        ## 得到更新后的损失
        new_loss = original_loss + 0.01 * addiontial_loss
        return new_loss
    
## 2. 把自己的 policy 封装在算法中： 
##    继承自PPO， 创建一个新的算法类， 默认调用的是自定义的policy
class MY_PPO(PPO):
    ## 重写 get_default_policy_class 函数， 使其返回自定义的policy 
    def get_default_policy_class(self, config):
        return MY_PPOTorchPolicy

## 三、使用自己的策略创建智能体，执行训练
## method-1
# config = PPOConfig(algo_class = MY_PPO) ## 配置使用自己的算法
# config = config.environment("CartPole-v1")
# config = config.rollouts(num_rollout_workers=2)
# config = config.framework(framework="torch") 
# algo = config.build()
# result = algo.train()
# print(pretty_print(result))

## method-2
config = PPOConfig()
config = config.environment("CartPole-v1")
config = config.rollouts(num_rollout_workers=2)
config = config.framework(framework="torch") 
algo = MY_PPO(config=config,)  ## 在这里配置使用自己的policy
result = algo.train()
print(pretty_print(result))

五、后记：

        如何既要修改网络结构，又要修改loss函数， 可以结合 上一篇博客 和本博客共同实现。