|
from typing import Union, Optional, List, Any, Tuple |
|
import os |
|
import torch |
|
import logging |
|
from functools import partial |
|
from tensorboardX import SummaryWriter |
|
from copy import deepcopy |
|
|
|
from ding.envs import get_vec_env_setting, create_env_manager |
|
from ding.worker import BaseLearner, InteractionSerialEvaluator, BaseSerialCommander, create_buffer, \ |
|
create_serial_collector |
|
from ding.config import read_config, compile_config |
|
from ding.policy import create_policy, PolicyFactory |
|
from ding.reward_model import create_reward_model |
|
from ding.utils import set_pkg_seed |
|
from ding.data.level_replay.level_sampler import LevelSampler |
|
from ding.policy.common_utils import default_preprocess_learn |
|
|
|
|
|
def generate_seeds(num_seeds=500, base_seed=0): |
|
return [base_seed + i for i in range(num_seeds)] |
|
|
|
|
|
def serial_pipeline_plr( |
|
input_cfg: Union[str, Tuple[dict, dict]], |
|
seed: int = 0, |
|
env_setting: Optional[List[Any]] = None, |
|
model: Optional[torch.nn.Module] = None, |
|
max_train_iter: Optional[int] = int(1e10), |
|
max_env_step: Optional[int] = int(1e10), |
|
) -> 'Policy': |
|
""" |
|
Overview: |
|
Serial pipeline entry for Priority Level Replay. |
|
Arguments: |
|
- input_cfg (:obj:`Union[str, Tuple[dict, dict]]`): Config in dict type. \ |
|
``str`` type means config file path. \ |
|
``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. |
|
- 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. |
|
""" |
|
if isinstance(input_cfg, str): |
|
cfg, create_cfg = read_config(input_cfg) |
|
else: |
|
cfg, create_cfg = deepcopy(input_cfg) |
|
create_cfg.policy.type = create_cfg.policy.type + '_command' |
|
env_fn = None if env_setting is None else env_setting[0] |
|
cfg = compile_config(cfg, seed=seed, env=env_fn, auto=True, create_cfg=create_cfg, save_cfg=True) |
|
collector_env_num = cfg.env.collector_env_num |
|
|
|
if env_setting is None: |
|
env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env) |
|
else: |
|
env_fn, collector_env_cfg, evaluator_env_cfg = env_setting |
|
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, dynamic_seed=False) |
|
evaluator_env.seed(cfg.seed, dynamic_seed=True) |
|
train_seeds = generate_seeds() |
|
level_sampler = LevelSampler( |
|
train_seeds, cfg.policy.model.obs_shape, cfg.policy.model.action_shape, collector_env_num, cfg.level_replay |
|
) |
|
set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) |
|
policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval', 'command']) |
|
|
|
|
|
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) |
|
collector = create_serial_collector( |
|
cfg.policy.collect.collector, |
|
env=collector_env, |
|
policy=policy.collect_mode, |
|
tb_logger=tb_logger, |
|
exp_name=cfg.exp_name |
|
) |
|
evaluator = InteractionSerialEvaluator( |
|
cfg.policy.eval.evaluator, evaluator_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name |
|
) |
|
commander = BaseSerialCommander( |
|
cfg.policy.other.commander, learner, collector, evaluator, None, policy.command_mode |
|
) |
|
|
|
|
|
|
|
|
|
|
|
learner.call_hook('before_run') |
|
|
|
seeds = [int(level_sampler.sample('sequential')) for _ in range(collector_env_num)] |
|
|
|
level_seeds = torch.Tensor(seeds) |
|
|
|
collector_env.seed(seeds) |
|
collector_env.reset() |
|
|
|
while True: |
|
collect_kwargs = commander.step() |
|
|
|
if evaluator.should_eval(learner.train_iter): |
|
stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) |
|
if stop: |
|
break |
|
|
|
new_data = collector.collect( |
|
train_iter=learner.train_iter, level_seeds=level_seeds, policy_kwargs=collect_kwargs |
|
) |
|
|
|
learner.train(new_data, collector.envstep) |
|
stacked_data = default_preprocess_learn(new_data, ignore_done=cfg.policy.learn.ignore_done, use_nstep=False) |
|
level_sampler.update_with_rollouts(stacked_data, collector_env_num) |
|
seeds = [int(level_sampler.sample()) for _ in range(collector_env_num)] |
|
level_seeds = torch.Tensor(seeds) |
|
collector_env.seed(seeds) |
|
collector_env.reset() |
|
if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: |
|
break |
|
|
|
learner.call_hook('after_run') |
|
return policy |
|
|