#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     52s
date:       2024/9/27
file:       train.py
version:
doc:
"""
import os
import sys
from argparse import ArgumentParser
from typing import List, Tuple, Type
os.environ["RAY_DEDUP_LOGS"] = "0"  # 完全显示ray训练过程中的日志信息
import ray
from ray.rllib.algorithms.registry import ALGORITHMS
from ray.train import CheckpointConfig

from corekit.callbacks.base_callbacks import BaseCallbacks
from corekit.config.base import BaseTaskConfig
from corekit.env import BaseEnv
from corekit.utils import WORK_DIR
from corekit.utils import get_debug_logger
from corekit.utils.results_collection import BaseResultsCollection
from corekit.utils.typedict.node import NodeInfo

_logger = get_debug_logger()
# 工程路径, ip地址
DEFAULAT_WORK_DIR = WORK_DIR
DEFAULT_IP_ADDRESS = '127.0.0.1'
os.environ['work_dir'] = DEFAULAT_WORK_DIR
os.environ['ip_address'] = DEFAULT_IP_ADDRESS


class TrainTaskConfig(BaseTaskConfig):
    env: Type["BaseEnv"] = None

    def __init__(self, engine: str = 'demo', mode: str = 'train', category: str = 'tune'):
        super().__init__(engine, mode, category)
        if self.engine == "LingYi":
            from envs.LingYi.env import LingYiEnv
            self.env = LingYiEnv
        assert self.env

    def start(self, multi_algo: bool = False, **kwargs) -> None:
        config = self.config
        assert config

        # 初始化ray
        self.init_ray(config)
        # 执行任务
        if multi_algo:  # 多算法测试
            algos = [
                'APPO', 'BC', 'CQL', 'DQN', 'DreamerV3', 'IMPALA',
                'MARWIL', 'PPO', 'RNNSAC', 'SAC', 'SimpleQ'
            ]

            done_list = []
            # done_list.extend(['A2C', 'A3C', 'BC', 'MARWIL', 'PG', 'PPO', 'IMPALA'])  # 直接成功
            # done_list.extend(['APEX', 'DQN', 'SimpleQ'])  # 仅支持 gym.spaces.Discrete
            # done_list.extend(['DDPG', 'TD3'])  # 改成torch, 当动作空间连续时, 将探索策略改为 EpsilonGreedy 后成功, 失败了?
            # done_list.extend(['ARS', 'ES'])  # 支持离散/多维离散动作空间, 将env.step()中的奖励设为 一个数而非MultiAgentDict 后成功
            # done_list.extend(['SAC'])  # 当动作空间离散时, 让模型继承SACTFModel后成功
            # # done_list.extend(['DDPPO'])  # 改成torch, 动作空间离散, 模型 TorchSharedWeightsModel, 其他改动见 settings.py 中的注释
            # # done_list.extend(['BanditLinTS', 'BanditLinUCB'])  # 改成torch, 将探索策略改为 ThompsonSampling 后成功
            # done_list.extend([
            #     'APEX_DDPG',
            # ])  # 改成torch, 当动作空间连续时, 将探索策略改为 EpsilonGreedy 后成功, 注在tf下会因为数据维度问题报错, 暂时不知道原因
            # done_list.extend([
            #     'R2D2', 'RNNSAC'
            # ])
            fail_list = []
            # fail_list.extend([
            #     'DREAMER',  # 改成torch, num_workers = 0, 连续动作空间, 内部好像有卷积过程, 需要特定的网络模型, 会有奇怪的数组维度问题, 无从解决
            #     'MBMPO',  # 改成torch, 'MultiAgentBatch' object has no attribute 'concat', 好像无解
            #     'QMIX',  # 改成torch, 状态和动作空间必须是 gym.spaces.Tuple, 且动作空间需由离散空间组成
            #     'RNNSAC'  # 改成torch和RNNSACTorchModel, mat1 and mat2 shapes cannot be multiplied (32x73 and 80x256)
            # ])  # 只能使用torch框架, 不能使用tf框架
            # fail_list.extend(['R2D2'])  # 需要调整模型为RNN模型, 目前的RNN模型不能用
            # fail_list.extend([
            #     'MAML'
            # ])  # 模型无关的元学习, 从元训练任务集中随机采样一批任务, 需要环境有 sample_tasks 函数
            # fail_list.extend(['CQL'])  # 将模型调整为SACTFModel后报错 'Categorical' object has no attribute 'sample_logp'
            # fail_list.extend([
            #     'SlateQ'
            # ])  # 动作空间要求MultiDiscrete, 观测空间必须能self.obs_space["doc"]["0"].shape[0], 必须使用 SlateQTFModel
            # fail_list.extend([
            #     'BanditLinTS', 'BanditLinUCB',
            # ])  # 会爆内存
            yes_list = []
            no_list = []
            debug_list = [
                # 支持的20种算法
                # 'A2C', 'A3C', 'APEX', 'APEX_DDPG', 'APPO', 'ARS', 'BC',
                # 'DDPG', 'DDPPO', 'DQN', 'ES', 'IMPALA', 'MARWIL', 'PG', 'PPO',
                # 'R2D2', 'RNNSAC', 'SAC', 'SimpleQ', 'TD3'
            ]

            debug_list.extend([
                # 'BanditLinTS', 'BanditLinUCB',  # 极慢, 甚至会爆内存
                # 'DREAMER',  # 内部好像有卷积过程, 需要特定的网络模型
                # 'QMIX',  # 多智能体强化学习算法,
                # 'MBMPO',  # 很可惜, sample()的返回结果是MultiAgentBatch而不是SampleBatch, 导致出错

                # 'ARS', 'ES',  # 成功, 但疯狂reset, 这2种算法不会执行on_trainer_init函数
                # 'APPO', 'DDPG',  # 有报错, 但能继续跑
                # File "C:\Users\lenovo\anaconda3\envs\ray13\lib\site-packages\ray\rllib\
                # execution\multi_gpu_learner_thread.py", line 215, in _step
                #     else batch.policy_batches[pid],
                # KeyError: 'policy_3'
                # 'DDPPO',  # 阻塞
                'A2C', 'A3C', 'APEX', 'APEX_DDPG', 'APPO', 'ARS', 'BC',
                'DDPG', 'DDPPO', 'DQN', 'ES', 'IMPALA', 'MARWIL', 'PG',
                'PPO', 'R2D2', 'RNNSAC', 'SAC', 'SimpleQ', 'TD3'

            ])  # 目前还需要debug的

            # 遍历所有支持的算法, 测试是否无误
            for algo in algos:
                if algo in done_list + fail_list:
                    continue
                if debug_list and algo not in debug_list:
                    continue
                config['alg'] = algo  # 更新算法
                config = self._valid_framework_and_exploration_by_alg(config)  # 更新算法参数

                try:
                    _logger.info(f'{algo} 算法开始训练')
                    self._start(config)
                    _logger.info(f'{algo} 算法已完成训练')
                    yes_list.append(algo)
                except Exception as _e:
                    _logger.error(f'{algo} 算法训练失败')
                    _logger.error(str(_e))
                    no_list.append(algo)
                    continue
                self._start(config)
        else:
            _logger.info(f'开始执行 {self.engine}-{self.mode}-{self.category}任务')
            self._start(config)

    @staticmethod
    def init_ray(config: dict):
        import ray
        # 配置Ray运行时环境
        if config['training_mode'] == "local":  # ray 的训练模式配置
            ray.init(local_mode=True

                     )  # 已弃用
            # print('可用硬件资源:', ray.available_resources())
        elif config['training_mode'] == "normal":
            ray.init()
        elif config['training_mode'] == "cluster":
            ray_address = f"ray://{config['ip']}:{config['port']}"
            _logger.info(f'尝试集群, ray集群主节点: {ray_address}')
            ray.init(
                ignore_reinit_error=True,
                runtime_env={"working_dir": WORK_DIR}  # TODO, 会将results中的ckpt和h5文件一起上传, 可优化
            )
            _logger.info(f'集群初始化成功')
        else:
            print("运行模式异常，请确认个人配置")

        if ray.is_initialized():
            nodes: List[NodeInfo] = ray.nodes()
            if len(nodes) > 1:  # 多节点模式, 即在ray cluster上执行任务
                _logger.info(f'当前集群下的节点信息: ')
                for node in nodes:
                    resource_info = node['Resources']
                    _logger.info(
                        f"节点 {node['NodeManagerAddress']}:{node['NodeManagerPort']}, ID: {node['NodeID']}, "
                        f"cpu: {resource_info['CPU']}, gpu: {resource_info.get('GPU', 0.)}, "
                        f"memory: {resource_info['memory']}"
                    )

                # @ray.remote
                # def add_sys_path():
                #     """将代码工程目录添加到sys.path, 避免worker节点不认识本地模块"""
                #     if WORK_DIR not in sys.path:
                #         print(f'无本地路径, 已添加')
                #         sys.path.append(WORK_DIR)  # 添加也没有用
                #     else:
                #         print(f'工作节点的sys.path如下')
                #         for i in sys.path:
                #             print(i)
                #         # for i in os.listdir(WORK_DIR):
                #         #     print(i)
                #
                #     print(f'{os.path.exists(os.path.join(WORK_DIR, "corekit"))=}')
                #
                # ray.get(add_sys_path.remote())

    @staticmethod
    def _start(config: dict):
        # 将初始参数解析并分割成 all_conf, trainer_conf, stop_conf
        all_conf, trainer_conf, stop_conf = TrainTaskConfig._splited_args(config)
        # 实验、检查点模型、训练信息保存路径配置
        run_or_experiment = all_conf['run_or_experiment']
        checkpoint_config = all_conf['checkpoint_config']
        results_dir = all_conf['results_dir']

        from ray.tune import run
        results = ray.tune.run(
            run_or_experiment,
            config=trainer_conf,
            stop=stop_conf,
            storage_path=results_dir,
            checkpoint_score_attr="episode_reward_mean",  # 根据平均奖励选择最佳
            checkpoint_config=checkpoint_config
        )  # 224.649s / 5 iter
        # 2025-07-18 add
        best_trial = results.get_best_trial(metric="episode_reward_mean", mode="max")
        best_checkpoint_dir = results.get_best_checkpoint(best_trial, metric="episode_reward_mean", mode="max")
        if best_checkpoint_dir is not None:
            best_checkpoint_dir_path = best_checkpoint_dir.path
            _logger.info(f"Best checkpoint directory: {best_checkpoint_dir_path}")
        else:
            _logger.info("Best checkpoint directory not exist!!!")
        # add end
        ray.shutdown()
        return results

    @staticmethod
    def _splited_args(args: dict) -> Tuple[dict, dict, dict]:
        """
        读取解析参数字典, 分割成3个字典 ray_conf, trainer_conf, stop_conf.

        :param args: dict, 所有参数的字典
        :return: ray_conf, env_conf, stop_conf: dict, dict, dict
        """
        alg: str = args['alg']

        if alg == 'TD3':
            _, default_trainer_config = ALGORITHMS['DDPG']()
            trainer, _ = ALGORITHMS['TD3']()
        else:
            trainer, default_trainer_config = ALGORITHMS[alg]()
        # from ray.tune.registry import get_trainable_cls
        # trainer = get_trainable_cls(alg)  # 这个函数内部会有warning, 日志很丑
        default_trainer_config = trainer.get_default_config()  # 获取训练器默认参数配置
        # Trainer所需的参数
        trainer_conf: dict = {key: args[key] for key in args if key in default_trainer_config}
        # 其他参数
        custom_conf: dict = {key: args[key] for key in args if key not in default_trainer_config}
        trainer_conf['env_config'] = custom_conf  # 回调中会使用

        # 所有的参数设置
        args['run_or_experiment'] = trainer
        args['checkpoint_config'] = CheckpointConfig(**args['checkpoint_config'])

        # 停止条件参数配置
        stop_conf = {
            "training_iteration": args['stop_iters'],  # int
            "timesteps_total": args['stop_timesteps'],  # int
            "episode_reward_mean": args['stop_reward'],  # float
        }

        return args, trainer_conf, stop_conf

    @staticmethod
    def _valid_framework_and_exploration_by_alg(args: dict) -> dict:
        """检查/调整/更新/适配 20种学习算法的设置

        不同的算法对参数的设置要求不一, 主要设置可分为5个方面:
        0) 参数要求; 1) 框架要求; 2) 动作/状态空间要求; 3) 探索策略要求; 4) 模型要求;

        本函数仅调整 1, 3, 4, 5方面, 2)的调整仅给出空间类型,
        """
        # 1) 框架要求的知识, 只能使用torch框架, 无法在tf框架下运行
        framework_demand: set = {
            # ray限制必须使用torch框架
            'DDPPO', 'DREAMER', 'MBMPO', 'QMIX', 'RNNSAC', 'R2D2', 'BanditLinTS', 'BanditLinUCB',
            # ray无框架限制, 但tf框架下没跑通, 只能在torch下运行
            'APEX_DDPG', 'DDPG', 'TD3', 'APPO', 'MAML', 'SAC'  #
        }
        # 2) 动作/状态空间要求知识
        act_space_demand: dict = {
            'MultiDiscrete': {
                'A2C', 'A3C', 'BC', 'MARWIL', 'PG', 'PPO', 'APPO', 'IMPALA',
                'ARS', 'ES',  # 'DREAMER'
                'BanditLinTS', 'BanditLinUCB',  # 其他动作空间会轻松爆内存
                # 'QMIX'
                # 'CQL'
                # 'MBMPO',
                'MAML',
            },  # 是否支持多维离散空间的算法
            'Discrete': {
                'A2C', 'A3C', 'BC', 'MARWIL', 'PG', 'PPO', 'APPO', 'IMPALA', 'APEX',
                'DQN', 'SimpleQ', 'ARS', 'ES', 'SAC', 'DDPPO', 'R2D2',
                # 'BanditLinTS', 'BanditLinUCB',
                'MBMPO',
                'RNNSAC',
            },  # 是否支持 Discrete 空间
            'Box': {
                'A2C', 'A3C', 'BC', 'MARWIL', 'PG', 'PPO', 'APPO', 'IMPALA',
                'BanditLinTS', 'BanditLinUCB', 'APEX_DDPG', 'DDPG', 'TD3',  # 'DREAMER'
            },  # 是否支持 Box 空间
        }
        # 3) 探索策略要求知识
        # from ray.rllib.agents.mbmpo.utils import MBMPOExploration
        exploration_demand: dict = {
            'default': 'StochasticSampling',
            'APEX_DDPG': 'StochasticSampling',
            'DDPG': 'StochasticSampling',
            'TD3': 'GaussianNoise',
            'BanditLinTS': 'UpperConfidenceBound',
            'BanditLinUCB': 'ThompsonSampling',
            # 'MBMPO': MBMPOExploration
        }  # 可选的探索策略 StochasticSampling, EpsilonGreedy, ThompsonSampling

        if args['engine'] == 'missile':
            if args['alg'] != 'MADDPG':
                args['alg'] = 'MADDPG'
                _logger.warning(f'missile 只进行 MADDPG 算法训练, 已自动修正')

        alg: str = args['alg']

        # 3) 探索策略调整
        exploration_config = {"type": exploration_demand.get(alg, 'StochasticSampling')}
        args['exploration_config'] = exploration_config

        # 0) 参数要求调整
        if alg == 'PPO':
            # 调整与否都不影响能否跑通
            args['num_sgd_iter'] = 60
            args['lambda'] = 0.9
        elif alg == 'A3C':
            args['lambda'] = 0.9
        elif alg == 'DREAMER':
            args['num_workers'] = 0
        elif alg == 'DDPPO':
            args['rollout_fragment_length'] = args.pop('train_batch_size')
            args['num_gpus'] = 0
            args['policies_to_train'] = None
        elif alg in ['BanditLinTS', 'BanditLinUCB', ]:
            args['timesteps_per_iteration'] = 100
            args['rollout_fragment_length'] = 1
            args['train_batch_size'] = 1
        elif alg == 'CQL':
            args['num_gpus'] = max(args['num_gpus'], 1)
            # from ray.rllib.models.tf.tf_action_dist import TFActionDistribution
            # args['model'] = {'custom_action_dist': TFActionDistribution}
            args['normalize_actions'] = True
            args['_use_beta_distribution'] = False
            # args['input'] = 'd4rl'
        elif alg == 'MBMPO':
            args['num_gpus'] = max(args['num_gpus'], 1)
        elif alg == 'MADDPG':
            if args['framework'] not in ['tf', 'tf2']:
                args['framework'] = 'tf2'
                _logger.warning(f'当前算法 {alg} 只支持 tf 框架, 已自动将框架从torch调整为tf')
            args['agent_id'] = 0  # 必须是整数
            args['learning_starts'] = 800
            args['use_local_critic'] = True
            # args['use_state_preprocessor'] = True
        elif alg == 'ARS':
            # args['rollouts_used'] = 300
            # args['num_rollouts'] = 300
            args['num_workers'] = 2
            args['sgd_stepsize'] = 0.99
        elif alg == 'ES':
            # args['train_batch_size'] = 300
            args['episodes_per_batch'] = 1  # 每个批次跑多少局
        elif alg == 'BC':
            args['framework'] = 'tf2'
            args['enable_connectors'] = True  # 必须使用 connectors
            args['_disable_execution_plan_api'] = False
            args['_enable_new_api_stack'] = False
            args['exploration_config'] = {}
        elif alg == 'MARWIL':
            args['num_gpus'] = min(args['num_gpus'], 1)
            args['learning_starts'] = 100
            args['_disable_execution_plan_api'] = False
        elif alg == 'TD3':  # 失败的, 老问题 else batch.policy_batches[pid] KeyError: 'policy_3';新问题训练10次就会爆内存(16g)
            # 现有方案是不采用TD3Trainer, 直接采用DDPGTrainer
            args.update({
                "learning_starts": 1000,
                "actor_hiddens": [64, 64],
                "critic_hiddens": [32, 32],
            })
        # elif alg == 'DREAMER':  # 失败的
        #     # from ray.rllib.agents.dreamer.dreamer_model import DreamerModel
        #     from corekit.custom_ann import DreamerTorch as AnnModel
        #     # act_space = gym.spaces.Box(low=0, high=1, shape=(5,))  # 动作空间维度
        #     dreamer_model = {
        #         "custom_model": AnnModel,
        #         # RSSM/PlaNET parameters
        #         "deter_size": 1908,
        #         "stoch_size": 1,
        #         # CNN Decoder Encoder
        #         "depth_size": 1,
        #         # General Network Parameters
        #         "hidden_size": 64,
        #         # Action STD
        #         "action_init_std": 5.0,
        #     }
        #     args['dreamer_model'] = dreamer_model

        # 1) 框架调整
        if alg in framework_demand and args['framework'] != 'torch':
            args['framework'] = 'torch'
            _logger.warning(f'当前算法 {alg} 只支持 torch 框架, 已自动将框架从tf调整为torch')

        # 2) 动作 / 状态空间要求
        act_space: str = 'MultiDiscrete'
        obs_space: str = 'Box'
        if alg in act_space_demand['MultiDiscrete']:
            act_space = 'MultiDiscrete'
        elif alg in act_space_demand['Discrete']:
            act_space = 'Discrete'
        elif alg in act_space_demand['Box']:
            act_space = 'Box'
        elif alg == 'QMIX':
            act_space = 'Tuple'
            obs_space = 'Tuple'
        args['obs_space'] = obs_space
        args['act_space'] = act_space

        # 4) 模型调整
        framework: str = args['framework']
        if framework in ['tf']:
            if alg in ['SAC', 'CQL']:
                from corekit.dnn.custom_ann import SACTFModel as AnnModel
            else:
                from corekit.dnn.custom_ann import SharedWeightsModel1 as AnnModel
        elif framework in ['tf2']:
            if alg in ['SAC', 'CQL']:
                from corekit.dnn.custom_ann import SACTFModel as AnnModel
            else:
                from corekit.dnn.custom_ann import TF2SharedWeightsModel as AnnModel
        elif framework == 'torch':
            if alg in ['SAC', 'CQL']:
                from corekit.dnn.custom_ann import SACTorch as AnnModel
            elif alg in ['RNNSAC']:
                from corekit.dnn.custom_ann import RNNSACTorch as AnnModel
            elif alg in ['R2D2']:
                from corekit.dnn.custom_ann import RNNModelTorch as AnnModel
            else:
                from corekit.dnn.custom_ann import SharedWeightsModelTorch as AnnModel
        else:
            raise ValueError(f'不支持的框架{framework}')
        args['ann_model'] = AnnModel

        if args['framework'] == 'tf2':
            args['eager_tracing'] = False  # tf2 模式下必须, 否则会失败
        if args.get('input_'):
            args['explore'] = False  # 使用离线数据, 没必要探索

        return args

    def get_all_config(self):
        args: dict = {
            'as_test': True,
            'explore': True,
            'env_config': {},  # 这个字典可以随意添加数据, ray框架不会报错
            # 可选项, 部分算法采用默认值会使每一次训练的间隔特别长
            'learning_starts': 500,
            'target_network_update_freq': 100,
            'timesteps_per_iteration': 100,
            'min_time_s_per_reporting': 1
        }  # 初始化, 添加不可设置的参数

        # 可设置项更新, 设置优先级 命令行参数 > 前端参数 > 默认参数 ================================
        # 默认参数
        default_config = self.get_general_default_config()  # 通用参数的默认值
        # default_config.update(self.get_unique_default_config())  # 引擎独有参数的默认值
        args.update(default_config)
        # # 前端参数
        frontend_config = self.get_frontend_config()
        args.update(frontend_config)
        # 命令行参数
        command_args = self.get_command_config()
        for key, value in command_args.items():
            if value is None:
                continue
            args[key] = value
        # 可设置项更新, 设置优先级 命令行参数 > 前端参数 > 默认参数 ================================

        args = self.update_special_config(args)
        return args

    @classmethod
    def get_general_default_config(cls) -> dict:
        """获得训练任务通用配置的默认值"""
        general_config: dict = {
            'work_dir': DEFAULAT_WORK_DIR,
            'ip_address': DEFAULT_IP_ADDRESS,
            # 训练参数
            "alg": "PPO",
            "engine": "demo",
            # "connect": False,  # 是否需要使用通信服务
            # "ip": "127.0.0.1",  # 仅当 connect 为 True 时生效
            # "port": 5555,  # 仅当 connect 为 True 时生效
            # "protocol": "zmq",  # 仅当 connect 为 True 时生效
            "num_workers": 3,
            # "num_agents": 18,
            "train_batch_size": 100,
            "training_mode": "normal",  # local:本地训练, normal:本地训练(不可调试), cluster:集群训练, 需在配置文件中指定集群head节点ip,port
            "api": "tune",  # 测试分布式新增
            "framework": "tf",
            "lr": 5e-03,
            "no_done_at_end": False,
            "gamma": 0.95,
            "lambda": 0.9,
            "sgd_minibatch_size": 64,
            "store_model_interval": 1,
            "store_model_num": 300,
            "load_pretrained": False,
            "pretrained_path": os.path.join(DEFAULAT_WORK_DIR, 'results', 'demo', 'pretrained_model'),
            "rule_engine": False,
            "camp": "Red",
            "callbacks": BaseCallbacks,
            "results_collection": BaseResultsCollection(),  # 结果汇总实例化
            # 停止条件
            "stop_iters": 300,
            "stop_steps": 1000000,
            "stop_timesteps": 5000000,
            "stop_reward": 10000000,
            # 硬件资源
            "num_envs_per_worker": 1,
            "num_cpus_per_worker": 1,
            "num_gpus_per_worker": 0,
            "num_cpus": 10,
            "num_gpus": 0,
            # ckpt
            "checkpoint_config": {
                "num_to_keep": 3, "checkpoint_frequency": 1, "checkpoint_score_attribute": 'training_iteration'
            },
            # "enable_connectors": False,
            # "output": os.path.join(DEFAULAT_WORK_DIR, 'results', 'demo', 'offline'),
            # "input_": None,
            # "create_env_on_local_worker": True  # 当设置了离线数据源时, 将该参数设为 True, 甚至可以不再设置 env
            # 日志
            "log_path": os.path.join(WORK_DIR, "log")  # 日志路径
        }  # 训练默认设置
        return general_config

    @classmethod
    def get_unique_default_config(cls) -> dict:
        """获得训练任务独有配置的默认值"""
        return {}

    @classmethod
    def get_command_config(cls, alg: str = 'PPO') -> dict:
        parser = cls.add_general_command_config(alg)
        # parser = cls.add_unique_command_config(parser)
        command_parse = parser.parse_args()
        commandt_args = vars(command_parse)
        return commandt_args

    @classmethod
    def add_general_command_config(
        cls, alg: str = 'PPO', parser: ArgumentParser = None
    ) -> ArgumentParser:
        """为训练任务通用配置增加命令行配置功能
        :param alg: str, 算法名
        :param parser: argparse.ArgumentParser
        :return:
        """
        if parser is None:
            parser = ArgumentParser()
        # 外协意见, ray 多机多卡模式时, os.getcwd()可能会不可靠, '127.0.0.1'不可用, 应设置一个全局变量来保证工作路径和ip正确
        parser.add_argument("--work_dir", type=str, help='工程路径', default=None)
        parser.add_argument("--ip_address", type=str, help='ip地址', default=None)
        # 训练参数
        parser.add_argument("--alg", type=str, help='强化学习训练算法', default=alg)
        parser.add_argument("--connect", type=bool, help='是否需要通信服务', default=None)
        parser.add_argument("--ip", type=str, help=f'ip地址, 仅当 connect 为 True 时生效', default=None)
        parser.add_argument("--port", type=int, help=f'通信端口, 仅当 connect 为 True 时生效', default=None)
        parser.add_argument("--protocol", type=str, help=f'通信协议, 仅当 connect 为 True 时生效',
                            choices=['zmq', 'udp', 'grpc'], default='zmq' if 'win' in sys.platform else None)
        parser.add_argument("--num_workers", type=int, help='训练的worker数量', default=None)
        parser.add_argument("--num_agents", type=int, help='红方智能体个数', default=None)
        parser.add_argument("--train_batch_size", type=int, help='训练批次大小',
                            default=None if 'win' in sys.platform else None)
        # local: 本地训练(可调试), normal: 本地训练(不可调试), cluster: 集群训练, 需在配置文件中指定集群head节点ip, port.
        parser.add_argument("--training_mode", type=str, help='ray训练模式设置',
                            choices=['local', 'normal', 'cluster'], default=None)
        parser.add_argument("--api", type=str, help='使用哪个API进行ray分布式训练',
                            choices=['tune', 'built', 'train'], default=None)
        parser.add_argument("--framework", type=str, help='深度学习框架',
                            choices=['tf', 'tf2', 'torch'], default=None)
        parser.add_argument("--lr", type=float, help='学习率', default=None)
        parser.add_argument("--no_done_at_end", help='是否终止', type=bool, default=None)
        parser.add_argument("--gamma", type=float, help='折扣因子', default=None)
        parser.add_argument("--lambda", type=float, help='GAE参数', default=None)
        parser.add_argument("--sgd_minibatch_size", type=int, help='随机梯度最小批次大小', default=None)
        parser.add_argument("--store_model_interval", type=int, help='每训练多少次保存一次模型', default=None)
        parser.add_argument("--store_model_num", type=int, help='一共保存最新多少个模型', default=None)
        parser.add_argument("--load_pretrained", type=bool, help='是否加载预训练模型', default=None)
        parser.add_argument("--pretrained_path", type=str, help='预训练模型加载路径', default=None)
        parser.add_argument("--rule_engine", type=bool, help='是否采用规则引擎', default=None)
        parser.add_argument("--camp", type=str, help='待训练智能体所处的阵营',
                            choices=['red', 'blue'], default=None)
        # 停止条件
        parser.add_argument("--stop_iters", type=int, help='训练终止的迭代次数',
                            default=None if 'win' in sys.platform else None)
        parser.add_argument("--stop_steps", type=int, help='每一局最大采样步数', default=None)
        parser.add_argument("--stop_timesteps", type=int, help='训练完成最大采样步数', default=None)
        parser.add_argument("--stop_reward", type=float, help='训练完成最大奖励值', default=None)
        # 硬件资源
        parser.add_argument("--num_envs_per_worker", type=int, help='每个进程的环境数', default=None)
        parser.add_argument("--num_cpus_per_worker", type=int, help='每个进程的CPU数量', default=None)
        parser.add_argument("--num_gpus_per_worker", type=int, help='每个进程的GPU数量', default=None)
        parser.add_argument("--num_cpus", type=int, help='硬件设备的CPU数量', default=None)
        parser.add_argument("--num_gpus", type=int, help='硬件设备的GPU数量', default=None)
        # 日志路径
        parser.add_argument("--log_path", type=str, help='日志路径', default=None)
        return parser

    @classmethod
    def add_unique_command_config(cls, parser: ArgumentParser) -> ArgumentParser:
        """为训练任务独有配置添加命令行配置功能"""
        return parser

    def update_special_config(self, config: dict) -> dict:
        config['env'] = self.env
        # 特殊参数的补充/调整-----------------------------------------------------------------
        engine: str = self.engine
        work_dir: str = config['work_dir']
        # 环境变量的更新/设置
        os.environ['work_dir'] = config['work_dir']  # TODO, Linux下只在当前进程及子进程中有效, ray在非local模式下无法使用
        os.environ['ip_address'] = config['ip_address']
        # 决定 ray.init()中的模式
        config['local_mode']: bool = 'local' in config['training_mode']
        # 使用ray.tune.run训练时产生的训练结果保存路径, ray框架默认路径为用户目录下的 ~/ray_results
        config['results_dir'] = os.path.join(work_dir, 'results', engine, 'train_results')
        if not os.path.exists(config['results_dir']):
            os.makedirs(config['results_dir'])
        # ckpt 和 model 的保存路径
        # 20250711修改checkpoint路径
        config['pretrained_path'] = os.path.join(work_dir, 'results', engine)
        # config['pretrained_path'] = os.path.join(work_dir, 'envs', engine, 'agents')

        config['model_path'] = os.path.join(work_dir, 'results', engine, 'trained_model')
        # 特殊参数的补充/调整-----------------------------------------------------------------

        # 根据框架和算法, 检查/调整/更新参数字典args
        config: dict = self._valid_framework_and_exploration_by_alg(config)
        # 更新引擎训练所需的参数
        config = self.update_train_config(config)
        return config

    @classmethod
    def update_train_config(cls, args: dict) -> dict:
        """接口, 更新引擎训练相关的设置"""
        args.update(cls.get_space_config(args))
        args.update(cls.get_callback_config(args))
        args.update(cls.get_multiagent_config(args))
        args.update(cls.get_train_config(args))
        args.update(cls.get_rule_config(args))
        args.update(cls.get_server_config(args))
        cls.get_alg_config(args)
        return args

    @classmethod
    def get_space_config(cls, args: dict) -> dict:
        """获得引擎所需的动作/状态空间设置

        不同引擎在使用不同训练算法时所需的动作/状态空间不同, 只能由引擎自己设置
        """
        return {}

    @classmethod
    def get_rule_config(cls, args: dict) -> dict:
        """配置规则智能体参数"""
        return {}

    @classmethod
    def get_callback_config(cls, args: dict) -> dict:
        """获得引擎的回调器设置"""
        return {}

    @classmethod
    def get_multiagent_config(cls, args: dict) -> dict:
        """获得引擎的多智能体设置"""
        return {}

    @classmethod
    def get_server_config(cls, args:dict) -> dict:
        """返回服务器配置参数"""
        return {}
