Source code for lzero.entry.train_alphazero

import logging
import os
from functools import partial
from typing import Optional, Tuple

import torch
from ding.config import compile_config
from ding.envs import create_env_manager
from ding.envs import get_vec_env_setting
from ding.policy import create_policy
from ding.utils import set_pkg_seed
from ding.worker import BaseLearner, create_buffer
from tensorboardX import SummaryWriter

from lzero.policy import visit_count_temperature
from lzero.worker import AlphaZeroCollector, AlphaZeroEvaluator


[docs]def train_alphazero( input_cfg: Tuple[dict, dict], seed: int = 0, model: Optional[torch.nn.Module] = None, model_path: Optional[str] = None, max_train_iter: Optional[int] = int(1e10), max_env_step: Optional[int] = int(1e10), ) -> 'Policy': # noqa """ Overview: The train entry for AlphaZero. Arguments: - input_cfg (:obj:`Tuple[dict, dict]`): Config in dict type. ``Tuple[dict, dict]`` type means [user_config, create_cfg]. - seed (:obj:`int`): Random seed. - env_setting (:obj:`Optional[List[Any]]`): A list with 3 elements: \ ``BaseEnv`` subclass, collector env config, and evaluator env config. - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. - model_path (:obj:`Optional[str]`): The pretrained model path, which should point to the ckpt file of the pretrained model, and an absolute path is recommended. In LightZero, the path is usually something like ``exp_name/ckpt/ckpt_best.pth.tar``. - max_train_iter (:obj:`Optional[int]`): Maximum policy update iterations in training. - max_env_step (:obj:`Optional[int]`): Maximum collected environment interaction steps. Returns: - policy (:obj:`Policy`): Converged policy. """ cfg, create_cfg = input_cfg create_cfg.policy.type = create_cfg.policy.type if cfg.policy.cuda and torch.cuda.is_available(): cfg.policy.device = 'cuda' else: cfg.policy.device = 'cpu' cfg = compile_config(cfg, seed=seed, env=None, auto=True, create_cfg=create_cfg, save_cfg=True) # Create main components: env, policy env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) collector_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in collector_env_cfg]) evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) collector_env.seed(cfg.seed) evaluator_env.seed(cfg.seed, dynamic_seed=False) set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval']) # load pretrained model if model_path is not None: policy.learn_mode.load_state_dict(torch.load(model_path, map_location=cfg.policy.device)) # Create worker components: learner, collector, evaluator, replay buffer, commander. tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) replay_buffer = create_buffer(cfg.policy.other.replay_buffer, tb_logger=tb_logger, exp_name=cfg.exp_name) policy_config = cfg.policy batch_size = policy_config.batch_size collector = AlphaZeroCollector( env=collector_env, policy=policy.collect_mode, tb_logger=tb_logger, exp_name=cfg.exp_name, ) evaluator = AlphaZeroEvaluator( eval_freq=cfg.policy.eval_freq, n_evaluator_episode=cfg.env.n_evaluator_episode, stop_value=cfg.env.stop_value, env=evaluator_env, policy=policy.eval_mode, tb_logger=tb_logger, exp_name=cfg.exp_name, ) # ============================================================== # Main loop # ============================================================== # Learner's before_run hook. learner.call_hook('before_run') if cfg.policy.update_per_collect is not None: update_per_collect = cfg.policy.update_per_collect while True: collect_kwargs = {} # set temperature for visit count distributions according to the train_iter, # please refer to Appendix D in MuZero paper for details. collect_kwargs['temperature'] = visit_count_temperature( policy_config.manual_temperature_decay, policy_config.fixed_temperature_value, policy_config.threshold_training_steps_for_final_temperature, trained_steps=learner.train_iter ) # Evaluate policy performance if evaluator.should_eval(learner.train_iter) and learner.train_iter > 0: stop, reward = evaluator.eval( learner.save_checkpoint, learner.train_iter, collector.envstep, ) if stop: break # Collect data by default config n_sample/n_episode new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs) new_data = sum(new_data, []) if cfg.policy.update_per_collect is None: # update_per_collect is None, then update_per_collect is set to the number of collected transitions multiplied by the replay_ratio. collected_transitions_num = len(new_data) update_per_collect = int(collected_transitions_num * cfg.policy.replay_ratio) replay_buffer.push(new_data, cur_collector_envstep=collector.envstep) # Learn policy from collected data for i in range(update_per_collect): # Learner will train ``update_per_collect`` times in one iteration. train_data = replay_buffer.sample(batch_size, learner.train_iter) if train_data is None: logging.warning( 'The data in replay_buffer is not sufficient to sample a mini-batch.' 'continue to collect now ....' ) break learner.train(train_data, collector.envstep) if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: break # Learner's after_run hook. learner.call_hook('after_run') return policy