#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/7/11
file:       base.py
version:
doc:
"""
import abc
from typing import Dict, Optional

from ray.rllib import SampleBatch
from ray.rllib.policy.policy import Policy
# from ray.rllib.utils.deprecation import DEPRECATED_VALUE
from ray.rllib.utils.typing import ModelWeights, TensorType


class BasePolicy(Policy, abc.ABC):
    """Policy 是 RLlib 中所有策略实现的基础抽象类

    它是所有特定于深度学习框架的子类(例如 TFPolicy 或 TorchPolicy)的超类.该类提供了以下功能的接口:

    1. 从观察值(及可能的其他输入)中计算动作.这通常涉及到根据当前策略和观察到的环境状态决定采取的动作.

    2. 管理策略的神经网络模型, 包括导出和加载模型权重.这使得策略能够保存和恢复其学习到的状态, 以便在不同的会话或环境中使用.

    3. 后处理从环境或其他输入获得的轨迹, 通过 postprocess_trajectory 方法.这可能包括添加额外的信息,
       如奖励的折扣累积, 或进行必要的数据转换, 以便后续的训练或分析.

    4. 从训练批次中计算损失.这是评估策略性能的关键步骤, 通常涉及比较策略的预测与实际结果, 并据此调整模型参数.

    5. 从训练批次中对神经网络模型执行更新, 这通常包括损失计算, 可以以两种方式完成：

        a. 在一个单一的步骤中进行(learn_on_batch).
        这种方式将整个批次的数据一次性送入模型, 计算损失并更新参数.

        b. 通过预加载批次数据, 然后进行n步的实际损失计算和更新(load_batch_into_buffer + learn_on_loaded_batch).
        这种分步处理的方式允许更细粒度的控制和可能的更高效利用计算资源."""

    is_train: bool = True  # 该策略是否进行训练

    def learn_on_batch(self, samples: SampleBatch) -> Dict[str, TensorType]:
        if not self.is_train:  # 非训练策略
            return {}

    def get_weights(self) -> ModelWeights:
        if not self.is_train:  # 非训练策略
            return {}

    def set_weights(self, weights: ModelWeights) -> None:
        if not self.is_train:  # 非训练策略
            return

    def export_model(self, export_dir: str, onnx: Optional[int] = None) -> None:
        if not self.is_train:  # 非训练策略
            return


class BaseRulePolicy(BasePolicy, abc.ABC):
    """规则策略基类"""
    is_train: bool = False
