DI-engine/ding/policy/bcq.py

from typing import List, Dict, Any, Tuple, Union
from collections import namedtuple
import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from ding.torch_utils import Adam, to_device
from ding.rl_utils import v_1step_td_data, v_1step_td_error, get_train_sample, get_nstep_return_data
from ding.model import model_wrap
from ding.policy import Policy
from ding.utils import POLICY_REGISTRY
from ding.utils.data import default_collate, default_decollate
from .common_utils import default_preprocess_learn


@POLICY_REGISTRY.register('bcq')
class BCQPolicy(Policy):
    config = dict(
        type='bcq',
        # (bool) Whether to use cuda for network.
        cuda=False,
        # (bool type) priority: Determine whether to use priority in buffer sample.
        # Default False in SAC.
        priority=False,
        # (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True.
        priority_IS_weight=False,
        # (int) Number of training samples(randomly collected) in replay buffer when training starts.
        # Default 10000 in SAC.
        random_collect_size=10000,
        nstep=1,
        model=dict(
            # (List) Hidden list for actor network head.
            actor_head_hidden_size=[400, 300],

            # (List) Hidden list for critic network head.
            critic_head_hidden_size=[400, 300],
            # Max perturbation hyper-parameter for BCQ
            phi=0.05,
        ),
        learn=dict(

            # How many updates(iterations) to train after collector's one collection.
            # Bigger "update_per_collect" means bigger off-policy.
            # collect data -> update policy-> collect data -> ...
            update_per_collect=1,
            # (int) Minibatch size for gradient descent.
            batch_size=100,

            # (float type) learning_rate_q: Learning rate for soft q network.
            # Default to 3e-4.
            # Please set to 1e-3, when model.value_network is True.
            learning_rate_q=3e-4,
            # (float type) learning_rate_policy: Learning rate for policy network.
            # Default to 3e-4.
            # Please set to 1e-3, when model.value_network is True.
            learning_rate_policy=3e-4,
            # (float type) learning_rate_vae: Learning rate for vae network.
            # `learning_rate_value` should be initialized, when model.vae_network is True.
            # Please set to 3e-4, when model.vae_network is True.
            learning_rate_vae=3e-4,
            # (bool) Whether ignore done(usually for max step termination env. e.g. pendulum)
            # Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers.
            # These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000.
            # However, interaction with HalfCheetah always gets done with done is False,
            # Since we inplace done==True with done==False to keep
            # TD-error accurate computation(``gamma * (1 - done) * next_v + reward``),
            # when the episode step is greater than max episode step.
            ignore_done=False,

            # (float type) target_theta: Used for soft update of the target network,
            # aka. Interpolation factor in polyak averaging for target networks.
            # Default to 0.005.
            target_theta=0.005,
            # (float) discount factor for the discounted sum of rewards, aka. gamma.
            discount_factor=0.99,
            lmbda=0.75,

            # (float) Weight uniform initialization range in the last output layer
            init_w=3e-3,
        ),
        collect=dict(
            # (int) Cut trajectories into pieces with length "unroll_len".
            unroll_len=1,
        ),
        eval=dict(),
        other=dict(
            replay_buffer=dict(
                # (int type) replay_buffer_size: Max size of replay buffer.
                replay_buffer_size=1000000,
                # (int type) max_use: Max use times of one data in the buffer.
                # Data will be removed once used for too many times.
                # Default to infinite.
                # max_use=256,
            ),
        ),
    )

    def default_model(self) -> Tuple[str, List[str]]:
        return 'bcq', ['ding.model.template.bcq']

    def _init_learn(self) -> None:
        r"""
        Overview:
            Learn mode init method. Called by ``self.__init__``.
            Init q, value and policy's optimizers, algorithm config, main and target models.
        """
        # Init
        self._priority = self._cfg.priority
        self._priority_IS_weight = self._cfg.priority_IS_weight
        self.lmbda = self._cfg.learn.lmbda
        self.latent_dim = self._cfg.model.action_shape * 2

        # Optimizers
        self._optimizer_q = Adam(
            self._model.critic.parameters(),
            lr=self._cfg.learn.learning_rate_q,
        )
        self._optimizer_policy = Adam(
            self._model.actor.parameters(),
            lr=self._cfg.learn.learning_rate_policy,
        )
        self._optimizer_vae = Adam(
            self._model.vae.parameters(),
            lr=self._cfg.learn.learning_rate_vae,
        )

        # Algorithm config
        self._gamma = self._cfg.learn.discount_factor

        # 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}
        )
        self._learn_model = model_wrap(self._model, wrapper_name='base')
        self._learn_model.reset()
        self._target_model.reset()

        self._forward_learn_cnt = 0

    def _forward_learn(self, data: dict) -> Dict[str, Any]:
        loss_dict = {}

        data = default_preprocess_learn(
            data,
            use_priority=self._priority,
            use_priority_IS_weight=self._cfg.priority_IS_weight,
            ignore_done=self._cfg.learn.ignore_done,
            use_nstep=False
        )
        if len(data.get('action').shape) == 1:
            data['action'] = data['action'].reshape(-1, 1)

        if self._cuda:
            data = to_device(data, self._device)

        self._learn_model.train()
        self._target_model.train()
        obs = data['obs']
        next_obs = data['next_obs']
        reward = data['reward']
        done = data['done']
        batch_size = obs.shape[0]

        # train_vae
        vae_out = self._model.forward(data, mode='compute_vae')
        recon, mean, log_std = vae_out['recons_action'], vae_out['mu'], vae_out['log_var']
        recons_loss = F.mse_loss(recon, data['action'])
        kld_loss = torch.mean(-0.5 * torch.sum(1 + log_std - mean ** 2 - log_std.exp(), dim=1), dim=0)
        loss_dict['recons_loss'] = recons_loss
        loss_dict['kld_loss'] = kld_loss
        vae_loss = recons_loss + 0.5 * kld_loss
        loss_dict['vae_loss'] = vae_loss
        self._optimizer_vae.zero_grad()
        vae_loss.backward()
        self._optimizer_vae.step()

        # train_critic
        q_value = self._learn_model.forward(data, mode='compute_critic')['q_value']

        with torch.no_grad():
            next_obs_rep = torch.repeat_interleave(next_obs, 10, 0)
            z = torch.randn((next_obs_rep.shape[0], self.latent_dim)).to(self._device).clamp(-0.5, 0.5)
            vae_action = self._model.vae.decode_with_obs(z, next_obs_rep)['reconstruction_action']
            next_action = self._target_model.forward({
                'obs': next_obs_rep,
                'action': vae_action
            }, mode='compute_actor')['action']

            next_data = {'obs': next_obs_rep, 'action': next_action}
            target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value']
            # the value of a policy according to the maximum entropy objective
            # find min one as target q value
            target_q_value = self.lmbda * torch.min(target_q_value[0], target_q_value[1]) \
                + (1 - self.lmbda) * torch.max(target_q_value[0], target_q_value[1])
            target_q_value = target_q_value.reshape(batch_size, -1).max(1)[0].reshape(-1, 1)

        q_data0 = v_1step_td_data(q_value[0], target_q_value, reward, done, data['weight'])
        loss_dict['critic_loss'], td_error_per_sample0 = v_1step_td_error(q_data0, self._gamma)
        q_data1 = v_1step_td_data(q_value[1], target_q_value, reward, done, data['weight'])
        loss_dict['twin_critic_loss'], td_error_per_sample1 = v_1step_td_error(q_data1, self._gamma)
        td_error_per_sample = (td_error_per_sample0 + td_error_per_sample1) / 2

        self._optimizer_q.zero_grad()
        (loss_dict['critic_loss'] + loss_dict['twin_critic_loss']).backward()
        self._optimizer_q.step()

        # train_policy
        z = torch.randn((obs.shape[0], self.latent_dim)).to(self._device).clamp(-0.5, 0.5)
        sample_action = self._model.vae.decode_with_obs(z, obs)['reconstruction_action']
        input = {'obs': obs, 'action': sample_action}
        perturbed_action = self._model.forward(input, mode='compute_actor')['action']
        q_input = {'obs': obs, 'action': perturbed_action}
        q = self._learn_model.forward(q_input, mode='compute_critic')['q_value'][0]
        loss_dict['actor_loss'] = -q.mean()
        self._optimizer_policy.zero_grad()
        loss_dict['actor_loss'].backward()
        self._optimizer_policy.step()
        self._forward_learn_cnt += 1
        self._target_model.update(self._learn_model.state_dict())
        return {
            'td_error': td_error_per_sample.detach().mean().item(),
            'target_q_value': target_q_value.detach().mean().item(),
            **loss_dict
        }

    def _monitor_vars_learn(self) -> List[str]:
        return [
            'td_error', 'target_q_value', 'critic_loss', 'twin_critic_loss', 'actor_loss', 'recons_loss', 'kld_loss',
            'vae_loss'
        ]

    def _state_dict_learn(self) -> Dict[str, Any]:
        ret = {
            'model': self._learn_model.state_dict(),
            'target_model': self._target_model.state_dict(),
            'optimizer_q': self._optimizer_q.state_dict(),
            'optimizer_policy': self._optimizer_policy.state_dict(),
            'optimizer_vae': self._optimizer_vae.state_dict(),
        }
        return ret

    def _init_eval(self):
        self._eval_model = model_wrap(self._model, wrapper_name='base')
        self._eval_model.reset()

    def _forward_eval(self, data: dict) -> Dict[str, Any]:
        data_id = list(data.keys())
        data = default_collate(list(data.values()))
        if self._cuda:
            data = to_device(data, self._device)
        data = {'obs': data}
        self._eval_model.eval()
        with torch.no_grad():
            output = self._eval_model.forward(data, mode='compute_eval')
        if self._cuda:
            output = to_device(output, 'cpu')
        output = default_decollate(output)
        return {i: d for i, d in zip(data_id, output)}

    def _init_collect(self) -> None:
        self._unroll_len = self._cfg.collect.unroll_len
        self._gamma = self._cfg.discount_factor  # necessary for parallel
        self._nstep = self._cfg.nstep  # necessary for parallel
        self._collect_model = model_wrap(self._model, wrapper_name='eps_greedy_sample')
        self._collect_model.reset()

    def _forward_collect(self, data: dict, **kwargs) -> dict:
        pass

    def _process_transition(self, obs: Any, model_output: dict, timestep: namedtuple) -> dict:
        pass

    def _get_train_sample(self, data: list) -> Union[None, List[Any]]:
        r"""
            Overview:
                Get the trajectory and the n step return data, then sample from the n_step return data
            Arguments:
                - data (:obj:`list`): The trajectory's cache
            Returns:
                - samples (:obj:`dict`): The training samples generated
            """
        data = get_nstep_return_data(data, self._nstep, gamma=self._gamma)
        return get_train_sample(data, self._unroll_len)