Please enable JavaScript.

Coggle requires JavaScript to display documents.

rllib, train - Coggle Diagram

- - - - 强化学习训练管理类：
        Trainer
        
        def init(self,config: TrainerConfigDict = None,env: str = None,logger_creator: Callable[[], Logger] = None):
        ---> 1. 注册环境：self._env_id = self._register_if_needed(env or config.get("env"))
        
        default_resource_request(cls, config: PartialTrainerConfigDict) -> Resources:
        --->1. cf = dict(cls._default_config, *config)
        --->2. Trainer._validate_config(cf)
        --->3. Resources(cpu=cf["num_cpus_for_driver"], gpu=cf["num_gpus"], memory=cf["memory"], object_store_memory=cf["object_store_memory"], extra_cpu=cf["num_cpus_per_worker"] num_workers, extra_gpu=cf["num_gpus_per_worker"] num_workers, extra_memory=cf["memory_per_worker"] num_workers, extra_object_store_memory=cf["object_store_memory_per_worker"] * num_workers)
        
        train(self) -> ResultDict:
        --->1. result = Trainable.train(self)
        --->2. 更新远程filter：self._sync_filters_if_needed(self.workers)
        --->3. 在“ evaluation_config”设置下评估当前策略：evaluation_metrics = self._evaluate()
        
        setup(self, config: PartialTrainerConfigDict):
        --->1.获取环境的creator
        --->2.初始化方法：self._init(self.config, self.env_creator)
        --->3.创建评估worker：self.evaluation_workers = self._make_workers(env_creator=self.env_creator,validate_env=None,policy_class=self._policy_class,config=merge_dicts(self.config, extra_config),num_workers=self.config["evaluation_num_workers"])
        
        cleanup(self):
        --->1.self.workers.stop()
        --->2.self.optimizer.stop()
        
        save_checkpoint(self, checkpoint_dir: str) -> str:
        --->1. pickle.dump(self.getstate(), open(checkpoint_path, "wb"))
        
        load_checkpoint(self, checkpoint_path: str):
        --->1. extra_data = pickle.load(open(checkpoint_path, "rb"))
        
        _make_workers(env_creator=self.env_creator,validate_env=None,policy_class=self._policy_class,config=merge_dicts(self.config, extra_config),num_workers=self.config["evaluation_num_workers"])：
        
        创建worker的工厂方法
        return WorkerSet(env_creator=env_creator,validate_env=validate_env,
        policy_class=policy_class,trainer_config=config,
        num_workers=num_workers,logdir=self.logdir)
        
        compute_action(self, observation: TensorStructType, state: List[TensorStructType] = None,prev_action: TensorStructType = None, prev_reward: float = None, info: EnvInfoDict = None,policy_id: PolicyID = DEFAULT_POLICY_ID, full_fetch: bool = False,explore: bool = None) -> TensorStructType:
        --->在本地worker上，指定policy计算一个动作
        
        _sync_weights_to_workers(self,*,worker_set: Optional[WorkerSet] = None,workers: Optional[List[RolloutWorker]] = None,) -> None:
        保存主worker中的权重到其它worker中
        weights = ray.put(self.workers.local_worker().save())
        worker_set.foreach_worker(lambda w: w.restore(ray.get(weights)))
        
        get_policy(self, policy_id: PolicyID = DEFAULT_POLICY_ID) -> Policy:
        
        get_weights(self, policies: List[PolicyID] = None) -> dict:
        
        set_weights(self, weights: Dict[PolicyID, dict]):
        
        _try_recover(self):
        尝试找出并删除任何有问题的worker。
        
        trainer构建方法
        
        trainer_template.py
        
        def build_trainer(name: str,*,default_config: Optional[TrainerConfigDict] = None,validate_config: Optional[Callable[[TrainerConfigDict], None]] = None,default_policy: Optional[Type[Policy]] = None,get_policy_class: Optional[Callable[[TrainerConfigDict], Optional[Type[Policy]]]] = None,validate_env: Optional[Callable[[EnvType, EnvContext], None]] = None,before_init: Optional[Callable[[Trainer], None]] = None,after_init: Optional[Callable[[Trainer], None]] = None,before_evaluate_fn: Optional[Callable[[Trainer], None]] = None,mixins: Optional[List[type]] = None,execution_plan: Optional[Callable[[WorkerSet, TrainerConfigDict], Iterable[ResultDict]]] = default_execution_plan) -> Type[Trainer]:
        --》name：表示trainer的名字，例如“PPO”
        --》default_config：强化学习算法的配置
        --》validate_config：用来验证配置的准确性
        --》default_policy：默认策略类
        --》get_policy_class：处理配置并返回策略类
        --》validate_env：去验证生成的环境，只在worker=0使用
        --》before_init：在构造worker之前执行，trainer实例作为参数
        --》after_init：在构造worker之后执行，trainer实例作为参数
        --》before_evaluate_fn：在评估前运行，trainer实例作为参数
        --》mixins：课程训练的trainer类别
        --》execution_plan：设置分布式训练工作流
        
        def build_trainer
        
        class trainer_cls(Trainer)
        
        def_ _init__(self, config=None, env=None, logger_creator=None):
        
        def _init(self, config: TrainerConfigDict,env_creator: Callable[[EnvConfigDict], EnvType]):
        --》创建worker：self.workers = self._make_workers(env_creator=env_creator,validate_env=validate_env,policy_class=self._policy_class,config=config,num_workers=self.config["num_workers"])
        
        def step(self):
        --->res = next(self.train_exec_impl)
        --->return res
        
        def getstate(self):
        
        def setstate(self, state):
        
        def with_updates(**overrides) -> Type[Trainer]:
        --》使用指定的overrides构建此Trainer类的副本。
  - - - torch版本的：
        TorchDistributionWrapper
        
        init(self, inputs: List[TensorType], model: ModelV2):
        self.last_sample = None
      - TensorFlow版本的：
        TFActionDistribution
        
        init(self, inputs: List[TensorType], model: ModelV2):
        self.sample_op = self._build_sample_op()
        self.sampled_action_logp_op = self.logp(self.sample_op)
        
        _build_sample_op(self) -> TensorType:
        取代sample()方法，
        
        实现类
        
        离散动作分类：
        Categorical
        
        多离散动作分类：
        MultiCategorical
        
        GumbelSoftmax
        
        基于高斯分布动作分类：
        DiagGaussian
        
        tanh-压缩版高斯分布的动作分类：
        SquashedGaussian
        
        对一组动作进行操作的动作分布：
        MultiActionDistribution
      - JAXDistribution
      - BinaryAutoregressiveDistribution
  - - - RemoteVectorEnv
        
        def poll(self)
        
        def send_actions(self, action_dict: MultiEnvDict) -> None:
        
        def try_reset(self,env_id: Optional[EnvID] = None) -> Optional[MultiAgentDict]:
        
        def stop(self) -> None:
      - _ExternalEnvToBaseEnv
      - _VectorEnvToBaseEnv
      - _MultiAgentEnvToBaseEnv
- - - - 1. 创建trainer：
        PPOTrainer = build_trainer(name="PPO",
        default_config=DEFAULT_CONFIG,
        validate_config=validate_config,
        default_policy=PPOTFPolicy,
        get_policy_class=get_policy_class,
        execution_plan=execution_plan,)
      - 2.默认配置：
        DEFAULT_CONFIG
      - 3.验证trainer的配置：
        def validate_config(config: TrainerConfigDict) -> None:
      - 4.PPO的策略类：
        PPOTFPolicy
      - 5.PPO算法的执行计划。定义分布式数据流：
        def execution_plan(workers: WorkerSet,config: TrainerConfigDict) -> LocalIterator[dict]:
        --->5.1.并行采样迭代器：rollouts = ParallelRollouts(workers, mode="bulk_sync")
        --->5.2.迭代更新policy并更新workers：
        train_op = rollouts.for_each(TrainOneStep(workers,num_sgd_iter=config["num_sgd_iter"],sgd_minibatch_size=config["sgd_minibatch_size"]))
        --->5.3.每轮训练后更新KL：
        train_op = train_op.for_each(lambda t: t[1]).for_each(UpdateKL(workers))
        --->5.4.返回评估指标：
        return StandardMetricsReporting(train_op, workers, config).for_each(lambda result: warn_about_bad_reward_scales(config, result))
- - - - RolloutWorker：通用经验采集器
        此类包装了一个策略实例和一个环境类，以从环境中收集经验。您可以创建许多此类的副本作为Ray演员来扩展RL训练。
        
        def _ _init__():
        
        def sample(self) -> SampleBatchType:
        
        def sample_with_count(self) -> Tuple[SampleBatchType, int]:
        
        def get_weights(self, policies: List[PolicyID] = None) -> (ModelWeights, dict):
        
        def set_weights(self, weights: ModelWeights,global_vars: dict = None) -> None:
        
        def compute_gradients(self, samples: SampleBatchType) -> Tuple[ModelGradients, dict]:
        
        def apply_gradients(self, grads: ModelGradients) -> Dict[PolicyID, Any]:
        
        def learn_on_batch(self, samples: SampleBatchType) -> dict:
        
        def sample_and_learn(self, expected_batch_size: int, num_sgd_iter: int, sgd_minibatch_size: str,standardize_fields: List[str]) -> Tuple[dict, int]:
        
        def as_remote(cls, num_cpus: int = None, num_gpus: int = None, memory: int = None, object_store_memory: int = None, resources: dict = None) -> type: