`ding.policy.atoc`¶

`ding.policy.atoc` ¶

`ATOCPolicy` ¶

Bases: Policy
Overview
Policy class of ATOC algorithm.
Interface: init, set_setting, repr, state_dict_handle Property: learn_mode, collect_mode, eval_mode
Full Source Code

../ding/policy/atoc.py
from typing import List, Dict, Any, Tuple, Unionfrom collections import namedtupleimport copyimport torchfrom ding.torch_utils import Adam, to_devicefrom ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_samplefrom ding.model import model_wrapfrom ding.utils import POLICY_REGISTRYfrom ding.utils.data import default_collate, default_decollatefrom .base_policy import Policyfrom .common_utils import default_preprocess_learn@POLICY_REGISTRY.register('atoc')class ATOCPolicy(Policy):    r"""    Overview:        Policy class of ATOC algorithm.    Interface:        __init__, set_setting, __repr__, state_dict_handle    Property:        learn_mode, collect_mode, eval_mode    """    config = dict(        # (str) RL policy register name (refer to function "POLICY_REGISTRY").        type='atoc',        # (bool) Whether to use cuda for network.        cuda=False,        # (bool) whether use on-policy training pipeline(behaviour policy and training policy are the same)        on_policy=False,        # (bool) Whether use priority(priority sample, IS weight, update priority)        priority=False,        # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True.        priority_IS_weight=False,        model=dict(            # (bool) Whether to use communication module in ATOC, if not, it is a multi-agent DDPG            communication=True,            # (int) The number of thought size            thought_size=8,            # (int) The number of agent for each communication group            agent_per_group=2,        ),        learn=dict(            # (int) Collect n_sample data, update model n_iteration time            update_per_collect=5,            # (int) The number of data for a train iteration            batch_size=64,            # (float) Gradient-descent step size of actor            learning_rate_actor=0.001,            # (float) Gradient-descent step size of critic            learning_rate_critic=0.001,            # ==============================================================            # The following configs is algorithm-specific            # ==============================================================            # (float) Target network update weight, theta * new_w + (1 - theta) * old_w, defaults in [0, 0.1]            target_theta=0.005,            # (float) Discount factor for future reward, defaults int [0, 1]            discount_factor=0.99,            # (bool) Whether to use communication module in ATOC, if not, it is a multi-agent DDPG            communication=True,            # (int) The frequency of actor update, each critic update            actor_update_freq=1,            # (bool) Whether use noise in action output when learning            noise=True,            # (float) The std of noise distribution for target policy smooth            noise_sigma=0.15,            # (float, float) The minimum and maximum value of noise            noise_range=dict(                min=-0.5,                max=0.5,            ),            # (bool) Whether to use reward batch norm in the total batch            reward_batch_norm=False,            ignore_done=False,        ),        collect=dict(            # (int) Collect n_sample data, update model n_iteration time            # n_sample=64,            # (int) Unroll length of a train iteration(gradient update step)            unroll_len=1,            # ==============================================================            # The following configs is algorithm-specific            # ==============================================================            # (float) The std of noise distribution for exploration            noise_sigma=0.4,        ),        eval=dict(),        other=dict(            replay_buffer=dict(                # (int) The max size of replay buffer                replay_buffer_size=100000,                # (int) The max use count of data, if count is bigger than this value, the data will be removed                max_use=10,            ),        ),    )    def default_model(self) -> Tuple[str, List[str]]:        return 'atoc', ['ding.model.template.atoc']    def _init_learn(self) -> None:        r"""        Overview:            Learn mode init method. Called by ``self.__init__``.            Init actor and critic optimizers, algorithm config, main and target models.        """        self._priority = self._cfg.priority        self._priority_IS_weight = self._cfg.priority_IS_weight        assert not self._priority and not self._priority_IS_weight        # algorithm config        self._communication = self._cfg.learn.communication        self._gamma = self._cfg.learn.discount_factor        self._actor_update_freq = self._cfg.learn.actor_update_freq        # actor and critic optimizer        self._optimizer_actor = Adam(            self._model.actor.parameters(),            lr=self._cfg.learn.learning_rate_actor,        )        self._optimizer_critic = Adam(            self._model.critic.parameters(),            lr=self._cfg.learn.learning_rate_critic,        )        if self._communication:            self._optimizer_actor_attention = Adam(                self._model.actor.attention.parameters(),                lr=self._cfg.learn.learning_rate_actor,            )        self._reward_batch_norm = self._cfg.learn.reward_batch_norm        # main and target models        self._target_model = copy.deepcopy(self._model)        self._target_model = model_wrap(            self._target_model,            wrapper_name='target',            update_type='momentum',            update_kwargs={'theta': self._cfg.learn.target_theta}        )        if self._cfg.learn.noise:            self._target_model = model_wrap(                self._target_model,                wrapper_name='action_noise',                noise_type='gauss',                noise_kwargs={                    'mu': 0.0,                    'sigma': self._cfg.learn.noise_sigma                },                noise_range=self._cfg.learn.noise_range            )        self._learn_model = model_wrap(self._model, wrapper_name='base')        self._learn_model.reset()        self._target_model.reset()        self._forward_learn_cnt = 0  # count iterations    def _forward_learn(self, data: dict) -> Dict[str, Any]:        r"""        Overview:            Forward and backward function of learn mode.        Arguments:            - data (:obj:`dict`): Dict type data, including at least ['obs', 'action', 'reward', 'next_obs']        Returns:            - info_dict (:obj:`Dict[str, Any]`): Including at least actor and critic lr, different losses.        """        loss_dict = {}        data = default_preprocess_learn(data, ignore_done=self._cfg.learn.ignore_done, use_nstep=False)        if self._cuda:            data = to_device(data, self._device)        # ====================        # critic learn forward        # ====================        self._learn_model.train()        self._target_model.train()        next_obs = data['next_obs']        reward = data['reward']        if self._reward_batch_norm:            reward = (reward - reward.mean()) / (reward.std() + 1e-8)        # current q value        q_value = self._learn_model.forward(data, mode='compute_critic')['q_value']        # target q value.        with torch.no_grad():            next_action = self._target_model.forward(next_obs, mode='compute_actor')['action']            next_data = {'obs': next_obs, 'action': next_action}            target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value']        td_data = v_1step_td_data(q_value.mean(-1), target_q_value.mean(-1), reward, data['done'], data['weight'])        critic_loss, td_error_per_sample = v_1step_td_error(td_data, self._gamma)        loss_dict['critic_loss'] = critic_loss        # ================        # critic update        # ================        self._optimizer_critic.zero_grad()        critic_loss.backward()        self._optimizer_critic.step()        # ===============================        # actor learn forward and update        # ===============================        # actor updates every ``self._actor_update_freq`` iters        if (self._forward_learn_cnt + 1) % self._actor_update_freq == 0:            if self._communication:                output = self._learn_model.forward(data['obs'], mode='compute_actor', get_delta_q=False)                output['delta_q'] = data['delta_q']                attention_loss = self._learn_model.forward(output, mode='optimize_actor_attention')['loss']                loss_dict['attention_loss'] = attention_loss                self._optimizer_actor_attention.zero_grad()                attention_loss.backward()                self._optimizer_actor_attention.step()            output = self._learn_model.forward(data['obs'], mode='compute_actor', get_delta_q=False)            critic_input = {'obs': data['obs'], 'action': output['action']}            actor_loss = -self._learn_model.forward(critic_input, mode='compute_critic')['q_value'].mean()            loss_dict['actor_loss'] = actor_loss            # actor update            self._optimizer_actor.zero_grad()            actor_loss.backward()            self._optimizer_actor.step()        # =============        # after update        # =============        loss_dict['total_loss'] = sum(loss_dict.values())        self._forward_learn_cnt += 1        self._target_model.update(self._learn_model.state_dict())        return {            'cur_lr_actor': self._optimizer_actor.defaults['lr'],            'cur_lr_critic': self._optimizer_critic.defaults['lr'],            'priority': td_error_per_sample.abs().tolist(),            'q_value': q_value.mean().item(),            **loss_dict,        }    def _state_dict_learn(self) -> Dict[str, Any]:        return {            'model': self._learn_model.state_dict(),            'target_model': self._target_model.state_dict(),            'optimizer_actor': self._optimizer_actor.state_dict(),            'optimizer_critic': self._optimizer_critic.state_dict(),            'optimize_actor_attention': self._optimizer_actor_attention.state_dict(),        }    def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None:        self._learn_model.load_state_dict(state_dict['model'])        self._target_model.load_state_dict(state_dict['target_model'])        self._optimizer_actor.load_state_dict(state_dict['optimizer_actor'])        self._optimizer_critic.load_state_dict(state_dict['optimizer_critic'])        self._optimizer_actor_attention.load_state_dict(state_dict['optimize_actor_attention'])    def _init_collect(self) -> None:        r"""        Overview:            Collect mode init method. Called by ``self.__init__``.            Init traj and unroll length, collect model.        """        self._unroll_len = self._cfg.collect.unroll_len        # collect model        self._collect_model = model_wrap(            self._model,            wrapper_name='action_noise',            noise_type='gauss',            noise_kwargs={                'mu': 0.0,                'sigma': self._cfg.collect.noise_sigma            },            noise_range=None,  # no noise clip in actor        )        self._collect_model.reset()    def _forward_collect(self, data: dict) -> dict:        r"""        Overview:            Forward function of collect mode.        Arguments:            - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \                values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer.        Returns:            - output (:obj:`Dict[int, Any]`): Dict type data, including at least inferred action according to input obs.        ReturnsKeys            - necessary: ``action``        """        data_id = list(data.keys())        data = default_collate(list(data.values()))        if self._cuda:            data = to_device(data, self._device)        self._collect_model.eval()        with torch.no_grad():            output = self._collect_model.forward(data, mode='compute_actor', get_delta_q=True)        if self._cuda:            output = to_device(output, 'cpu')        output = default_decollate(output)        return {i: d for i, d in zip(data_id, output)}    def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> Dict[str, Any]:        r"""        Overview:            Generate dict type transition data from inputs.        Arguments:            - obs (:obj:`Any`): Env observation            - model_output (:obj:`dict`): Output of collect model, including at least ['action']            - timestep (:obj:`namedtuple`): Output after env step, including at least ['obs', 'reward', 'done'] \                (here 'obs' indicates obs after env step, i.e. next_obs).        Return:            - transition (:obj:`Dict[str, Any]`): Dict type transition data.        """        if self._communication:            transition = {                'obs': obs,                'next_obs': timestep.obs,                'action': model_output['action'],                'delta_q': model_output['delta_q'],                'reward': timestep.reward,                'done': timestep.done,            }        else:            transition = {                'obs': obs,                'next_obs': timestep.obs,                'action': model_output['action'],                'reward': timestep.reward,                'done': timestep.done,            }        return transition    def _get_train_sample(self, data: list) -> Union[None, List[Any]]:        if self._communication:            delta_q_batch = [d['delta_q'] for d in data]            delta_min = torch.stack(delta_q_batch).min()            delta_max = torch.stack(delta_q_batch).max()            for i in range(len(data)):                data[i]['delta_q'] = (data[i]['delta_q'] - delta_min) / (delta_max - delta_min + 1e-8)        return get_train_sample(data, self._unroll_len)    def _init_eval(self) -> None:        r"""        Overview:            Evaluate mode init method. Called by ``self.__init__``.            Init eval model. Unlike learn and collect model, eval model does not need noise.        """        self._eval_model = model_wrap(self._model, wrapper_name='base')        self._eval_model.reset()    def _forward_eval(self, data: dict) -> dict:        r"""        Overview:            Forward function of eval mode, similar to ``self._forward_collect``.        Arguments:            - data (:obj:`Dict[str, Any]`): Dict type data, stacked env data for predicting policy_output(action), \                values are torch.Tensor or np.ndarray or dict/list combinations, keys are env_id indicated by integer.        Returns:            - output (:obj:`Dict[int, Any]`): The dict of predicting action for the interaction with env.        ReturnsKeys            - necessary: ``action``        """        data_id = list(data.keys())        data = default_collate(list(data.values()))        if self._cuda:            data = to_device(data, self._device)        self._eval_model.eval()        with torch.no_grad():            output = self._eval_model.forward(data, mode='compute_actor')        if self._cuda:            output = to_device(output, 'cpu')        output = default_decollate(output)        return {i: d for i, d in zip(data_id, output)}    def _monitor_vars_learn(self) -> List[str]:        r"""        Overview:            Return variables' name if variables are to used in monitor.        Returns:            - vars (:obj:`List[str]`): Variables' name list.        """        return [            'cur_lr_actor',            'cur_lr_critic',            'critic_loss',            'actor_loss',            'attention_loss',            'total_loss',            'q_value',        ]
ding.policy.atoc¶

ding.policy.atoc ¶

ATOCPolicy ¶

Full Source Code

`ding.policy.atoc`¶

`ding.policy.atoc` ¶

`ATOCPolicy` ¶
