Spaces:
Sleeping
Sleeping
| from typing import Union, Optional, List, Any, Tuple | |
| import torch | |
| import os | |
| from functools import partial | |
| from tensorboardX import SummaryWriter | |
| from copy import deepcopy | |
| from ding.worker import BaseLearner, InteractionSerialEvaluator, BaseSerialCommander, create_buffer, \ | |
| get_buffer_cls, create_serial_collector | |
| from ding.world_model import WorldModel | |
| from ding.worker import IBuffer | |
| from ding.envs import get_vec_env_setting, create_env_manager | |
| from ding.config import read_config, compile_config | |
| from ding.utils import set_pkg_seed, deep_merge_dicts | |
| from ding.policy import create_policy | |
| from ding.world_model import create_world_model | |
| from ding.entry.utils import random_collect | |
| def mbrl_entry_setup( | |
| input_cfg: Union[str, Tuple[dict, dict]], | |
| seed: int = 0, | |
| env_setting: Optional[List[Any]] = None, | |
| model: Optional[torch.nn.Module] = None, | |
| ) -> Tuple: | |
| 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) | |
| 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) | |
| evaluator_env.seed(cfg.seed, dynamic_seed=False) | |
| set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) | |
| # create logger | |
| tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) | |
| # create world model | |
| world_model = create_world_model(cfg.world_model, env_fn(cfg.env), tb_logger) | |
| # create policy | |
| policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval', 'command']) | |
| # create worker | |
| 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.eval_mode, tb_logger, exp_name=cfg.exp_name | |
| ) | |
| env_buffer = create_buffer(cfg.policy.other.replay_buffer, tb_logger=tb_logger, exp_name=cfg.exp_name) | |
| commander = BaseSerialCommander( | |
| cfg.policy.other.commander, learner, collector, evaluator, env_buffer, policy.command_mode | |
| ) | |
| return ( | |
| cfg, | |
| policy, | |
| world_model, | |
| env_buffer, | |
| learner, | |
| collector, | |
| collector_env, | |
| evaluator, | |
| commander, | |
| tb_logger, | |
| ) | |
| def create_img_buffer( | |
| cfg: dict, input_cfg: Union[str, Tuple[dict, dict]], world_model: WorldModel, tb_logger: 'SummaryWriter' | |
| ) -> IBuffer: # noqa | |
| if isinstance(input_cfg, str): | |
| _, create_cfg = read_config(input_cfg) | |
| else: | |
| _, create_cfg = input_cfg | |
| img_buffer_cfg = cfg.world_model.other.imagination_buffer | |
| img_buffer_cfg.update(create_cfg.imagination_buffer) | |
| buffer_cls = get_buffer_cls(img_buffer_cfg) | |
| cfg.world_model.other.imagination_buffer.update(deep_merge_dicts(buffer_cls.default_config(), img_buffer_cfg)) | |
| if img_buffer_cfg.type == 'elastic': | |
| img_buffer_cfg.set_buffer_size = world_model.buffer_size_scheduler | |
| img_buffer = create_buffer(cfg.world_model.other.imagination_buffer, tb_logger=tb_logger, exp_name=cfg.exp_name) | |
| return img_buffer | |
| def serial_pipeline_dyna( | |
| 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': # noqa | |
| """ | |
| Overview: | |
| Serial pipeline entry for dyna-style model-based RL. | |
| 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. | |
| """ | |
| cfg, policy, world_model, env_buffer, learner, collector, collector_env, evaluator, commander, tb_logger = \ | |
| mbrl_entry_setup(input_cfg, seed, env_setting, model) | |
| img_buffer = create_img_buffer(cfg, input_cfg, world_model, tb_logger) | |
| learner.call_hook('before_run') | |
| if cfg.policy.get('random_collect_size', 0) > 0: | |
| random_collect(cfg.policy, policy, collector, collector_env, commander, env_buffer) | |
| while True: | |
| collect_kwargs = commander.step() | |
| # eval the policy | |
| if evaluator.should_eval(collector.envstep): | |
| stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) | |
| if stop: | |
| break | |
| # fill environment buffer | |
| data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs) | |
| env_buffer.push(data, cur_collector_envstep=collector.envstep) | |
| # eval&train world model and fill imagination buffer | |
| if world_model.should_eval(collector.envstep): | |
| world_model.eval(env_buffer, collector.envstep, learner.train_iter) | |
| if world_model.should_train(collector.envstep): | |
| world_model.train(env_buffer, collector.envstep, learner.train_iter) | |
| world_model.fill_img_buffer( | |
| policy.collect_mode, env_buffer, img_buffer, collector.envstep, learner.train_iter | |
| ) | |
| for i in range(cfg.policy.learn.update_per_collect): | |
| batch_size = learner.policy.get_attribute('batch_size') | |
| train_data = world_model.sample(env_buffer, img_buffer, batch_size, learner.train_iter) | |
| learner.train(train_data, collector.envstep) | |
| if cfg.policy.on_policy: | |
| # On-policy algorithm must clear the replay buffer. | |
| env_buffer.clear() | |
| img_buffer.clear() | |
| if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: | |
| break | |
| learner.call_hook('after_run') | |
| return policy | |
| def serial_pipeline_dream( | |
| 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': # noqa | |
| """ | |
| Overview: | |
| Serial pipeline entry for dreamer-style model-based RL. | |
| 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. | |
| """ | |
| cfg, policy, world_model, env_buffer, learner, collector, collector_env, evaluator, commander, tb_logger = \ | |
| mbrl_entry_setup(input_cfg, seed, env_setting, model) | |
| learner.call_hook('before_run') | |
| if cfg.policy.get('random_collect_size', 0) > 0: | |
| random_collect(cfg.policy, policy, collector, collector_env, commander, env_buffer) | |
| while True: | |
| collect_kwargs = commander.step() | |
| # eval the policy | |
| if evaluator.should_eval(collector.envstep): | |
| stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep) | |
| if stop: | |
| break | |
| # fill environment buffer | |
| data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs) | |
| env_buffer.push(data, cur_collector_envstep=collector.envstep) | |
| # eval&train world model and fill imagination buffer | |
| if world_model.should_eval(collector.envstep): | |
| world_model.eval(env_buffer, collector.envstep, learner.train_iter) | |
| if world_model.should_train(collector.envstep): | |
| world_model.train(env_buffer, collector.envstep, learner.train_iter) | |
| update_per_collect = cfg.policy.learn.update_per_collect // world_model.rollout_length_scheduler( | |
| collector.envstep | |
| ) | |
| update_per_collect = max(1, update_per_collect) | |
| for i in range(update_per_collect): | |
| batch_size = learner.policy.get_attribute('batch_size') | |
| train_data = env_buffer.sample(batch_size, learner.train_iter) | |
| # dreamer-style: use pure on-policy imagined rollout to train policy, | |
| # which depends on the current envstep to decide the rollout length | |
| learner.train( | |
| train_data, collector.envstep, policy_kwargs=dict(world_model=world_model, envstep=collector.envstep) | |
| ) | |
| if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: | |
| break | |
| learner.call_hook('after_run') | |
| return policy | |
| def serial_pipeline_dreamer( | |
| 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': # noqa | |
| """ | |
| Overview: | |
| Serial pipeline entry for dreamerv3. | |
| 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. | |
| """ | |
| cfg, policy, world_model, env_buffer, learner, collector, collector_env, evaluator, commander, tb_logger = \ | |
| mbrl_entry_setup(input_cfg, seed, env_setting, model) | |
| learner.call_hook('before_run') | |
| # prefill environment buffer | |
| if cfg.policy.get('random_collect_size', 0) > 0: | |
| cfg.policy.random_collect_size = cfg.policy.random_collect_size // cfg.policy.collect.unroll_len | |
| random_collect(cfg.policy, policy, collector, collector_env, commander, env_buffer) | |
| while True: | |
| collect_kwargs = commander.step() | |
| # eval the policy | |
| if evaluator.should_eval(collector.envstep): | |
| stop, reward = evaluator.eval( | |
| learner.save_checkpoint, | |
| learner.train_iter, | |
| collector.envstep, | |
| policy_kwargs=dict(world_model=world_model) | |
| ) | |
| if stop: | |
| break | |
| # train world model and fill imagination buffer | |
| steps = ( | |
| cfg.world_model.pretrain | |
| if world_model.should_pretrain() else int(world_model.should_train(collector.envstep)) | |
| ) | |
| for _ in range(steps): | |
| batch_size = learner.policy.get_attribute('batch_size') | |
| batch_length = cfg.policy.learn.batch_length | |
| post, context = world_model.train( | |
| env_buffer, collector.envstep, learner.train_iter, batch_size, batch_length | |
| ) | |
| start = post | |
| learner.train( | |
| start, collector.envstep, policy_kwargs=dict(world_model=world_model, envstep=collector.envstep) | |
| ) | |
| # fill environment buffer | |
| data = collector.collect( | |
| train_iter=learner.train_iter, | |
| policy_kwargs=dict(world_model=world_model, envstep=collector.envstep, **collect_kwargs) | |
| ) | |
| env_buffer.push(data, cur_collector_envstep=collector.envstep) | |
| if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter: | |
| break | |
| learner.call_hook('after_run') | |
| return policy | |