from typing import Dict, List import math import random import pickle import torch from ding.utils import REWARD_MODEL_REGISTRY from .base_reward_model import BaseRewardModel def collect_state_action_pairs(iterator): # concat state and action """ Overview: Concate state and action pairs from input iterator. Arguments: - iterator (:obj:`Iterable`): Iterables with at least ``obs`` and ``action`` tensor keys. Returns: - res (:obj:`Torch.tensor`): State and action pairs. """ res = [] for item in iterator: state = item['obs'] action = item['action'] # s_a = torch.cat([state, action.float()], dim=-1) res.append((state, action)) return res @REWARD_MODEL_REGISTRY.register('pwil') class PwilRewardModel(BaseRewardModel): """ Overview: The Pwil reward model class (https://arxiv.org/pdf/2006.04678.pdf) Interface: ``estimate``, ``train``, ``load_expert_data``, ``collect_data``, ``clear_date``, \ ``__init__``, ``_train``, ``_get_state_distance``, ``_get_action_distance`` Config: == ================== ===== ============= ======================================= ======================= ID Symbol Type Default Value Description Other(Shape) == ================== ===== ============= ======================================= ======================= 1 ``type`` str pwil | Reward model register name, refer | | to registry ``REWARD_MODEL_REGISTRY`` | 2 | ``expert_data_`` str expert_data. | Path to the expert dataset | Should be a '.pkl' | ``path`` .pkl | | file 3 | ``sample_size`` int 1000 | sample data from expert dataset | | with fixed size | 4 | ``alpha`` int 5 | factor alpha | 5 | ``beta`` int 5 | factor beta | 6 | ``s_size`` int 4 | state size | 7 | ``a_size`` int 2 | action size | 8 | ``clear_buffer`` int 1 | clear buffer per fixed iters | make sure replay ``_per_iters`` | buffer's data count | isn't too few. | (code work in entry) == ================== ===== ============= ======================================= ======================= Properties: - reward_table (:obj: `Dict`): In this algorithm, reward model is a dictionary. """ config = dict( # (str) Reward model register name, refer to registry ``REWARD_MODEL_REGISTRY``. type='pwil', # (str) Path to the expert dataset. # expert_data_path='expert_data.pkl', # (int) Sample data from expert dataset with fixed size. sample_size=1000, # r = alpha * exp((-beta*T/sqrt(|s_size|+ |a_size|))*c_i) # key idea for this reward is to minimize. # the Wasserstein distance between the state-action distribution. # (int) Factor alpha. alpha=5, # (int) Factor beta. beta=5, #（int）State size. # s_size=4, # (int) Action size. # a_size=2, # (int) Clear buffer per fixed iters. clear_buffer_per_iters=1, ) def __init__(self, config: Dict, device: str, tb_logger: 'SummaryWriter') -> None: # noqa """ Overview: Initialize ``self.`` See ``help(type(self))`` for accurate signature. Arguments: - cfg (:obj:`Dict`): Training config - device (:obj:`str`): Device usage, i.e. "cpu" or "cuda" - tb_logger (:obj:`str`): Logger, defaultly set as 'SummaryWriter' for model summary """ super(PwilRewardModel, self).__init__() self.cfg: Dict = config assert device in ["cpu", "cuda"] or "cuda" in device self.device = device self.expert_data: List[tuple] = [] self.train_data: List[tuple] = [] # In this algo, model is a dict self.reward_table: Dict = {} self.T: int = 0 self.load_expert_data() def load_expert_data(self) -> None: """ Overview: Getting the expert data from ``config['expert_data_path']`` attribute in self Effects: This is a side effect function which updates the expert data attribute (e.g. ``self.expert_data``); \ in this algorithm, also the ``self.expert_s``, ``self.expert_a`` for states and actions are updated. """ with open(self.cfg.expert_data_path, 'rb') as f: self.expert_data = pickle.load(f) print("the data size is:", len(self.expert_data)) sample_size = min(self.cfg.sample_size, len(self.expert_data)) self.expert_data = random.sample(self.expert_data, sample_size) self.expert_data = [(item['obs'], item['action']) for item in self.expert_data] self.expert_s, self.expert_a = list(zip(*self.expert_data)) print('the expert data demonstrations is:', len(self.expert_data)) def collect_data(self, data: list) -> None: """ Overview: Collecting training data formatted by ``fn:concat_state_action_pairs``. Arguments: - data (:obj:`list`): Raw training data (e.g. some form of states, actions, obs, etc) Effects: - This is a side effect function which updates the data attribute in ``self``; \ in this algorithm, also the ``s_size``, ``a_size`` for states and actions are updated in the \ attribute in ``self.cfg`` Dict; ``reward_factor`` also updated as ``collect_data`` called. """ self.train_data.extend(collect_state_action_pairs(data)) self.T = len(self.train_data) s_size = self.cfg.s_size a_size = self.cfg.a_size beta = self.cfg.beta self.reward_factor = -beta * self.T / math.sqrt(s_size + a_size) def train(self) -> None: """ Overview: Training the Pwil reward model. """ self._train(self.train_data) def estimate(self, data: list) -> List[Dict]: """ Overview: Estimate reward by rewriting the reward key in each row of the data. Arguments: - data (:obj:`list`): the list of data used for estimation, \ with at least ``obs`` and ``action`` keys. Effects: - This is a side effect function which updates the ``reward_table`` with ``(obs,action)`` \ tuples from input. """ # NOTE: deepcopy reward part of data is very important, # otherwise the reward of data in the replay buffer will be incorrectly modified. train_data_augmented = self.reward_deepcopy(data) for item in train_data_augmented: s = item['obs'] a = item['action'] if (s, a) in self.reward_table: item['reward'] = self.reward_table[(s, a)] else: # when (s, a) pair is not trained, set the reward value to default value(e.g.: 0) item['reward'] = torch.zeros_like(item['reward']) return train_data_augmented def _get_state_distance(self, s1: list, s2: list) -> torch.Tensor: """ Overview: Getting distances of states given 2 state lists. One single state \ is of shape ``torch.Size([n])`` (``n`` referred in in-code comments) Arguments: - s1 (:obj:`torch.Tensor list`): the 1st states' list of size M - s2 (:obj:`torch.Tensor list`): the 2nd states' list of size N Returns: - distance (:obj:`torch.Tensor`) Euclidean distance tensor of \ the state tensor lists, of size M x N. """ # Format the values in the tensors to be of float type s1 = torch.stack(s1).float() s2 = torch.stack(s2).float() M, N = s1.shape[0], s2.shape[0] # Automatically fill in length s1 = s1.view(M, -1) s2 = s2.view(N, -1) # Automatically fill in & format the tensor size to be (MxNxn) s1 = s1.unsqueeze(1).repeat(1, N, 1) s2 = s2.unsqueeze(0).repeat(M, 1, 1) # Return the distance tensor of size MxN return ((s1 - s2) ** 2).mean(dim=-1) def _get_action_distance(self, a1: list, a2: list) -> torch.Tensor: # TODO the metric of action distance maybe different from envs """ Overview: Getting distances of actions given 2 action lists. One single action \ is of shape ``torch.Size([n])`` (``n`` referred in in-code comments) Arguments: - a1 (:obj:`torch.Tensor list`): the 1st actions' list of size M - a2 (:obj:`torch.Tensor list`): the 2nd actions' list of size N Returns: - distance (:obj:`torch.Tensor`) Euclidean distance tensor of \ the action tensor lists, of size M x N. """ a1 = torch.stack(a1).float() a2 = torch.stack(a2).float() M, N = a1.shape[0], a2.shape[0] a1 = a1.view(M, -1) a2 = a2.view(N, -1) a1 = a1.unsqueeze(1).repeat(1, N, 1) a2 = a2.unsqueeze(0).repeat(M, 1, 1) return ((a1 - a2) ** 2).mean(dim=-1) def _train(self, data: list): """ Overview: Helper function for ``train``, find the min disctance ``s_e``, ``a_e``. Arguments: - data (:obj:`list`): Raw training data (e.g. some form of states, actions, obs, etc) Effects: - This is a side effect function which updates the ``reward_table`` attribute in ``self`` . """ batch_s, batch_a = list(zip(*data)) s_distance_matrix = self._get_state_distance(batch_s, self.expert_s) a_distance_matrix = self._get_action_distance(batch_a, self.expert_a) distance_matrix = s_distance_matrix + a_distance_matrix w_e_list = [1 / len(self.expert_data)] * len(self.expert_data) for i, item in enumerate(data): s, a = item w_pi = 1 / self.T c = 0 expert_data_idx = torch.arange(len(self.expert_data)).tolist() while w_pi > 0: selected_dist = distance_matrix[i, expert_data_idx] nearest_distance = selected_dist.min().item() nearest_index_selected = selected_dist.argmin().item() nearest_index = expert_data_idx[nearest_index_selected] if w_pi >= w_e_list[nearest_index]: c = c + nearest_distance * w_e_list[nearest_index] w_pi = w_pi - w_e_list[nearest_index] expert_data_idx.pop(nearest_index_selected) else: c = c + w_pi * nearest_distance w_e_list[nearest_index] = w_e_list[nearest_index] - w_pi w_pi = 0 reward = self.cfg.alpha * math.exp(self.reward_factor * c) self.reward_table[(s, a)] = torch.FloatTensor([reward]) def clear_data(self) -> None: """ Overview: Clearing training data. \ This is a side effect function which clears the data attribute in ``self`` """ self.train_data.clear()