actorcritic_net.py

import torch
import torch.nn as nn
import torch.autograd as autograd
from torch.autograd import Variable
import torch.nn.functional as F
import torch.nn.init as init
from torch.distributions.categorical import Categorical
from torch.distributions.normal import Normal
from torch.nn.utils.rnn import pad_sequence
import math
import numpy as np

from decoder_only_transformer import DecoderOnlyTransformer, DeTrConfig


def mlp(sizes, activation, output_activation=nn.Identity()):
    layers = []
    for j in range(len(sizes)-1):
        act = activation if j < len(sizes)-2 else output_activation
        layers += [nn.Linear(sizes[j], sizes[j+1]), act]
    return nn.Sequential(*layers)

LOG_STD_MAX = 2
LOG_STD_MIN = -20
EPSILON = 1E-6
class ActorNetwork(nn.Module):
    """Useful as a baseline in REINFORCE updates"""
    def __init__(self,
            embedding_dim,
            hidden_dim,
            actions_dim,
            max_rounds):
        super(ActorNetwork, self).__init__()
        
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.actions_dim = actions_dim
        self.max_rounds = max_rounds
        self.mask_net = DecoderOnlyTransformer(
            config = DeTrConfig(n_block=max_rounds, \
                                n_embd=embedding_dim, \
                                n_outputs=max_rounds+actions_dim*3),
                                # n_outputs=2+actions_dim*3),
            type = 'actor')

    def forward(self, inputs, decode_type):
        """
        Args:
            inputs: [horizon, block_size(23), n_embd]
        """ 
        assert len(inputs.shape) == 3
        horizon = inputs.shape[0]
        batch_size = 1
        inputs = inputs.reshape(1, -1, inputs.shape[-1]) # [batch_size, n_token, n_embd]
        # mask_logits = torch.reshape(self.mask_net(inputs, generate_mode='inference'), (batch_size, self.actions_dim, 3))
        output = self.mask_net(inputs, generate_mode='inference')
        
        mask_logits = torch.reshape(output[:, self.max_rounds:], (batch_size, self.actions_dim, 3))
        mask_dist = Categorical(logits=mask_logits)

        # decide max rounds by categorical distribution
        max_rounds_logits = output[:, :self.max_rounds]
        max_rounds_dist = Categorical(logits=max_rounds_logits)
        
        ''' 
        # decide max rounds by normal distribution
        trunc_mu = output[:, 0]
        trunc_log_std = output[:, 1]
        trunc_log_std = torch.clamp(trunc_log_std, LOG_STD_MIN, LOG_STD_MAX)
        trunc_std = torch.exp(trunc_log_std)
        trunc_dist = Normal(trunc_mu, trunc_std)
        '''

        # trunc_logits = output[:, 0:2] * torch.tensor([horizon/self.max_rounds, 1 - horizon/self.max_rounds]).to(output.device)
        # trunc_dist = Categorical(logits=mask_logits) 

        if decode_type == "stochastic":
            mask_vec = mask_dist.sample() - 1
            max_rounds = max_rounds_dist.sample() + 1
            # trunc_ratio = trunc_dist.sample() # decide max rounds by normal distribution
            # if_trunc = trunc_dist.sample()
        elif decode_type == "greedy":
            mask_vec = torch.argmax(mask_logits, dim=-1) - 1
            max_rounds = torch.argmax(max_rounds_logits, dim=-1) + 1
            # trunc_ratio = trunc_mu # decide max rounds by normal distribution

        # trunc_ratio = (torch.tanh(trunc_ratio) + 1)*0.5 # torch.sigmoid? # decide max rounds by normal distribution

        return (mask_vec.squeeze(), max_rounds.squeeze())

    def logprobs(self, inputs, mask_idxes, max_rounds): # trunc_ratio_vec
        """ Propagate inputs through the network
        Args: 
            inputs: batch_size (num of episodes) * [horizon, block_size(23), n_embd]
        """
        batch_size = len(inputs)
        n_embd = inputs[0].shape[-1]
        device = inputs[0].device
        episode_lengths = [input.shape[0] for input in inputs]

        # padding
        padded_inputs = pad_sequence(inputs, batch_first=True) # [batch_size, horizon, block_size(23), n_embd]
        horizon = padded_inputs.shape[1]
        padded_inputs = padded_inputs.reshape(batch_size, -1, n_embd) # [batch_size, n_token, n_embd]
        mask_idxes = [torch.from_numpy(np.array(mask_idxes[i])) for i in range(batch_size)]
        padded_mask_idxes = pad_sequence(mask_idxes, batch_first=True, padding_value=0).to(device) # [batch_size, horizon, len(SEPA_LIST)]
        # decide max rounds by categorical distribution
        max_rounds = [torch.from_numpy(np.array(max_rounds[i])) for i in range(batch_size)]
        padded_max_rounds = pad_sequence(max_rounds, batch_first=True, padding_value=1).squeeze(dim=-1).to(device)
        ''' 
        # decide max rounds by normal distribution
        trunc_ratio_vec = [torch.from_numpy(np.arctanh(2*np.clip(np.array(trunc_ratio_vec[i]), EPSILON, 1-EPSILON)-1)) for i in range(batch_size)] # a = 0.5*(tanh(u)+1), u = tanh^(-1)(2*a - 1) 
        padded_trunc_ratio_vec = pad_sequence(trunc_ratio_vec, batch_first=True, padding_value=0).squeeze(dim=-1).to(device) # [batch_size, horizon, len(SEPA_LIST)]
        '''

        # compute logprob of mask_idxes
        output = self.mask_net(padded_inputs, generate_mode='train') # [batch_size, horizon, 2+len(SEPA_LIST)*3]

        mask_logits = torch.reshape(output[:,:, self.max_rounds:], (batch_size, -1, self.actions_dim, 3)) # (batch_size, horizon, len(SEPA_LIST), 3)
        mask_dist = Categorical(logits=mask_logits) # (batch_size, horizon, len(SEPA_LIST))
        mask_logprob = torch.sum(mask_dist.log_prob(padded_mask_idxes+1), dim=-1)

        # decide max rounds by categorical distribution
        max_rounds_logits = output[:, :, :self.max_rounds]
        max_rounds_dist = Categorical(logits=max_rounds_logits)
        max_rounds_logprob = max_rounds_dist.log_prob(padded_max_rounds-1)

        ''' 
        # decide max rounds by normal distribution
        # compute logprob of trunc_ratio_vec

        # prob_scales = torch.arange(horizon)*(1.0/self.max_rounds)
        # trunc_logits = output[:, :, 0:2] * torch.stack((prob_scales, 1-prob_scales), dim=-1) # (batch_size, horizon, 2)
        # trunc_dist = Categorical(logits=trunc_logits) # (batch_size, horizon)
        trunc_mu = output[:, :, 0]
        trunc_log_std = output[:, :, 1]
        trunc_log_std = torch.clamp(trunc_log_std, LOG_STD_MIN, LOG_STD_MAX)
        trunc_std = torch.exp(trunc_log_std)
        trunc_dist = Normal(trunc_mu, trunc_std)
        trunc_logprob = trunc_dist.log_prob(padded_trunc_ratio_vec) # (batch_size, horizon) # log(u)
        trunc_logprob -= 2*(np.log(2) - padded_trunc_ratio_vec - F.softplus(-2*padded_trunc_ratio_vec)) # logprob = log(u)-log(1-tanh(u)^2)

        logprob = mask_logprob + trunc_logprob
        '''
        logprob = mask_logprob + max_rounds_logprob
        logprob = [logprob[i][:episode_lengths[i]] for i in range(batch_size)]
        
        return mask_dist.probs, max_rounds_dist.probs, logprob # trunc_dist.mean

class CriticNetwork(nn.Module):
    """Useful as a baseline in REINFORCE updates"""
    def __init__(self, embedding_dim, max_rounds):
        super(CriticNetwork, self).__init__()
        self.value_net = DecoderOnlyTransformer(
            config = DeTrConfig(n_block=max_rounds, \
                                n_embd=embedding_dim, \
                                n_outputs=1),
            type = 'critic'
        )
    def forward(self, inputs):
        batch_size = len(inputs)
        n_embd = inputs[0].shape[-1]
        episode_lengths = [input.shape[0] for input in inputs]
        padded_inputs = pad_sequence(inputs, batch_first=True) # [batch_size, horizon, block_size(23), n_embd]
        padded_inputs = padded_inputs.reshape(batch_size, -1, n_embd) # [batch_size, n_token, n_embd]
        x = self.value_net(padded_inputs, generate_mode='train').squeeze(-1) # (batch_size, horizon)

        x = [x[i][:episode_lengths[i]] for i in range(batch_size)]

        return x