<jupyter_start><jupyter_text>So we want to<jupyter_code>from dotmap import DotMap
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import time
import sys
import os
from tqdm import tqdm
from torch import optim
from torch.distributions import Categorical<jupyter_output><empty_output><jupyter_text>Scaled Dot-Product AttentionWhen we take a dot product of query and key values, the variance of the result of the dot product may be scaled by $d$, which is the dimension of the key/query vectors. To ensure that the variance of the dot product still remains the same regardless of vector length, the scaled dot-product attention scoring function is used.Assuming $\mathbf{Q} \in \mathbb{R}^{n \times d}$, keys $\mathbf{K}\in \mathbb{R}^{m \times d}$, values $\mathbf{V}\in \mathbb{R}^{m \times v}$ we do$$\text{softmax} \left(\frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{d}}\right)\mathbf{V} \quad \in \mathbb{R}^{n\times v}$$ As the paper suggests, we apply masking when we do not want certain nodes to be considered for training. It makes sense to only consider those nodes which are part of a feasible sequence. A feasible sequence is basically [current node, $v_i$, $v_j$, $v_n$ ]<jupyter_code>class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k):
        super(ScaledDotProductAttention, self).__init__()
        self.scale_factor = np.sqrt(d_k)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, q, k, v, attn_mask=None):
        # q: [b_size x len_q x d_k]
        # k: [b_size x len_k x d_k]
        # v: [b_size x len_v x d_v] note: (len_k == len_v)
        # Batch-Matrix Multiplication
        attn = torch.bmm(q, k.transpose(1, 2)) / self.scale_factor  # attn: [b_size x len_q x len_k]
        if attn_mask is not None:
        #    assert attn_mask.size() == attn.size()
            attn.data.masked_fill_(attn_mask==0, -1e32)

        attn = self.softmax(attn )
        outputs = torch.bmm(attn, v) # outputs: [b_size x len_q x d_v]
        return outputs, attn
class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model, d_ff):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, x):
        residual = x # inputs: [b_size x len_q x d_model]
        outputs = self.w_2(F.relu(self.w_1(x)))
        return self.layer_norm(residual + outputs)
class _MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super(_MultiHeadAttention, self).__init__()

        self.d_k = d_model // n_heads
        self.d_v = d_model // n_heads
        self.d_model = d_model
        self.n_heads = n_heads

        self.w_q = nn.Parameter(torch.FloatTensor(n_heads, d_model, self.d_k))
        self.w_k = nn.Parameter(torch.FloatTensor(n_heads, d_model, self.d_k))
        self.w_v = nn.Parameter(torch.FloatTensor(n_heads, d_model, self.d_v))

        self.attention = ScaledDotProductAttention(self.d_k)

    def forward(self, q, k, v, attn_mask=None, is_adj=True):
        (d_k, d_v, d_model, n_heads) = (self.d_k, self.d_v, self.d_model, self.n_heads)
        b_size = k.size(0)

        q_s = q.repeat(n_heads, 1, 1).view(n_heads, -1, d_model)  # [n_heads x b_size * len_q x d_model]
        k_s = k.repeat(n_heads, 1, 1).view(n_heads, -1, d_model)  # [n_heads x b_size * len_k x d_model]
        v_s = v.repeat(n_heads, 1, 1).view(n_heads, -1, d_model)  # [n_heads x b_size * len_v x d_model]

        q_s = torch.bmm(q_s, self.w_q).view(b_size * n_heads, -1, d_k)  # [b_size * n_heads x len_q x d_k]
        k_s = torch.bmm(k_s, self.w_k).view(b_size * n_heads, -1, d_k)  # [b_size * n_heads x len_k x d_k]
        v_s = torch.bmm(v_s, self.w_v).view(b_size * n_heads, -1, d_v)  # [b_size * n_heads x len_v x d_v]

        if attn_mask is not None:
            if is_adj:
                outputs, attn = self.attention(q_s, k_s, v_s, attn_mask=attn_mask.repeat(n_heads, 1, 1))
            else:
                outputs, attn = self.attention(q_s, k_s, v_s, attn_mask=attn_mask.unsqueeze(1).repeat(n_heads, 1, 1))
        else:
            outputs, attn = self.attention(q_s, k_s, v_s, attn_mask=None)

        return torch.split(outputs, b_size, dim=0), attn
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super(MultiHeadAttention, self).__init__()

        self.d_k = d_model // n_heads
        self.attention = _MultiHeadAttention(d_model, n_heads)
        self.proj = nn.Linear(n_heads * self.d_k, d_model)
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, q, k, v, attn_mask = None, is_adj = True):
        # q: [b_size x len_q x d_model]
        # k: [b_size x len_k x d_model]
        # v: [b_size x len_v x d_model] note (len_k == len_v)
        residual = q
        # outputs: a list of tensors of shape [b_size x len_q x d_v] (length: n_heads)
        outputs, attn = self.attention(q, k, v, attn_mask=attn_mask, is_adj=is_adj)
        # concatenate 'n_heads' multi-head attentions
        outputs = torch.cat(outputs, dim=-1)
        # project back to residual size, result_size = [b_size x len_q x d_model]
        outputs = self.proj(outputs)

        return self.layer_norm(residual + outputs), attn
class EncoderLayer(nn.Module):
    def __init__(self, d_model, d_ff, n_heads):
        super(EncoderLayer, self).__init__()

        self.enc_self_attn = MultiHeadAttention(d_model, n_heads)
        self.pos_ffn = PositionwiseFeedForward(d_model, d_ff)

    def forward(self, enc_inp, rec_enc_inp, self_attn_mask):
        enc_outputs, attn = self.enc_self_attn(enc_inp, rec_enc_inp, enc_inp, attn_mask=self_attn_mask)
        enc_outputs = self.pos_ffn(enc_outputs)

        return enc_outputs, attn
class Encoder(nn.Module):
    def __init__(self, features_dim, dfeatures_dim, hidden_size, args):
        super(Encoder, self).__init__()

        n_heads = args.n_heads # number of heads
        d_ff = args.ff_dim # feed_forward_hidden
        n_layers = args.n_layers # number of Layers

        self.L1 = nn.Linear(features_dim, hidden_size//2) # for static features
        self.L2 = nn.Linear(dfeatures_dim, hidden_size//2) # for dynamic features

        self.layers = nn.ModuleList([EncoderLayer(hidden_size, d_ff, n_heads) for _ in range(n_layers)])

    def forward(self, emb_inp, rec_inp, mask, dummy_arg):
        for layer in self.layers:
            emb_inp, _ = layer(emb_inp, rec_inp, mask)

        return emb_inp

class Decoder(nn.Module):
    def __init__(self, hidden_size):
        super(Decoder, self).__init__()

        self.W1 = nn.Linear(hidden_size, hidden_size, bias=False)
        self.W2 = nn.Linear(hidden_size, hidden_size)
        self.V = nn.Parameter(torch.zeros((hidden_size, 1), requires_grad=True))

        self.first_h_0 = nn.Parameter(torch.FloatTensor(1, hidden_size), requires_grad=True)
        self.first_h_0.data.uniform_(-(1. / math.sqrt(hidden_size)), 1. / math.sqrt(hidden_size))

        self.c0 = nn.Parameter(torch.FloatTensor( 1, hidden_size),requires_grad=True)
        self.c0.data.uniform_(-(1. / math.sqrt(hidden_size)), 1. / math.sqrt(hidden_size))

        self.hidden_0 = (self.first_h_0, self.c0)

        self.lstm = nn.LSTMCell(hidden_size, hidden_size)


    def forward(self, input, hidden, enc_outputs, mask):
        hidden = self.lstm(input, hidden)
        w1e = self.W1(enc_outputs)
        w2h = self.W2(hidden[0]).unsqueeze(1)
        u = torch.tanh(w1e + w2h)
        a = u.matmul(self.V)
        a = 10*torch.tanh(a).squeeze(2)

        policy = F.softmax(a + mask.float().log(), dim=1)

        return policy, hidden
class RecPointerNetwork(nn.Module):

    def __init__(self, features_dim, dfeatures_dim, hidden_dim, args):
        super(RecPointerNetwork, self).__init__()

        self.features_dim = features_dim
        self.dfeatures_dim = dfeatures_dim
        self.use_checkpoint = args.use_checkpoint
        self.hidden_dim = hidden_dim
        self.decoder = Decoder(hidden_dim)
        self.encoder = Encoder(features_dim, dfeatures_dim, hidden_dim, args)
        # see https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/11
        self.dummy_tensor = torch.ones(1, dtype=torch.float32, requires_grad=True)

        self._initialize_parameters()

    def _initialize_parameters(self):
        for name, param in self.named_parameters():
            if len(param.shape) > 1:
                nn.init.xavier_uniform_(param)

    def _load_model_weights(self, path_string, device):
        self.load_state_dict(torch.load(path_string, map_location=device))


    def forward(self, enc_inputs, enc_hidden, adj_mask, dec_input, dec_hidden, mask, first_step=False):
        policy, dec_hidden, enc_outputs = self._one_step(enc_inputs, enc_hidden, adj_mask, dec_input, dec_hidden, mask, first_step)
        return policy, dec_hidden, enc_outputs

    def _one_step(self, enc_inputs, enc_hidden, adj_mask, dec_input, dec_hidden, mask, first_step):
        if self.use_checkpoint:
            enc_outputs = checkpoint(self.encoder, enc_inputs, enc_hidden, adj_mask, self.dummy_tensor)
        else:
            enc_outputs = self.encoder(enc_inputs, enc_hidden, adj_mask, self.dummy_tensor)

        if first_step:
            return  None, None, enc_outputs
        else:
            policy, dec_hidden = self.decoder(dec_input, dec_hidden, enc_outputs, mask)
            return policy, dec_hidden, enc_outputs

    def sta_emb(self, sta_inp):
        return torch.tanh(self.encoder.L1(sta_inp))

    def dyn_emb(self, dyn_inp):
        return torch.tanh(self.encoder.L2(dyn_inp))<jupyter_output><empty_output><jupyter_text>Utils<jupyter_code># config.py in the original code
cf = dict(
        BENCHMARK_INSTANCES_PATH = './data/benchmark/',
        GENERATED_INSTANCES_PATH = './data/generated/',
        RESULTS_PATH = './results'
)
cf = DotMap(cf)
# problem_config.py in the original code
pcf = DotMap(dict(
# Indices of instance data
X_COORDINATE_IDX = 0,
Y_COORDINATE_IDX = 1,
VIS_DURATION_TIME_IDX = 2,
OPENING_TIME_WINDOW_IDX = 3,
CLOSING_TIME_WINDOW_IDX = 4,
REWARD_IDX = 5,
ARRIVAL_TIME_IDX = 6,

# For generating instances
SAMP_DAY_FRAC_INF = 4/24,
UB_T_INIT_FRAC = 15/24,
LB_T_MAX_FRAC = 12/24,
CORR_SCORE_STD = 10,

MULTIPLE_SCORE = 1.1,

X_MAX = 100. # max square length (X_MAX)
))
def test_n_vert_1(instance_data, instance_type):
    N_VERT_ROW = 0

    if instance_type=='Gavalas':
        N_VERT_COL = 3
        DATA_INIT_ROW = 1
    else:
        N_VERT_COL = 2
        DATA_INIT_ROW = 2

    n_vert = instance_data[N_VERT_ROW][N_VERT_COL]
    count_vert = len(instance_data)-(DATA_INIT_ROW+1)

    assert count_vert==n_vert, 'number of vertices doesnt match number of data rows'


def test_n_vert_2(instance_data, instance_type):
    N_VERT_ROW = 0
    if instance_type=='Gavalas':
        N_VERT_COL = 3
    else:
        N_VERT_COL = 2
    COLUMN_NAMES = ['vertex number', 'x coordinate', 'y coordinate',
                    'service duration or visiting time', 'profit of the location',
                    'not relevant 1', 'not relevant 2', 'not relevant 3',
                    'opening of time window', 'closing of time window']
    COLUMN_NAMES = [s.replace(' ', '_') for s in COLUMN_NAMES]

    VERTEX_NUMBER_COL = [i for i,n in enumerate(COLUMN_NAMES) if n=='vertex_number'][0]
    n_vert = instance_data[N_VERT_ROW][N_VERT_COL]
    last_vert_number = instance_data[-1][VERTEX_NUMBER_COL]

    assert last_vert_number==n_vert, 'number of vertices doesnt match vertice count of last row'


def test_n_vert_3(instance_data, instance_type):
    if instance_type=='Gavalas':
        N_DAYS_INDEX = 1
        n_days = int(np.array(instance_data[0])[N_DAYS_INDEX])
        assert n_days==1, 'not a single tour/1 day instance'
    else:
        pass
def parse_instance_data_Gavalas(instance_data):
    """parse instance data into dataframe"""

    # get start date
    N_DAYS_INDEX = 1
    START_DAY_INDEX = 2
    M = instance_data[0][N_DAYS_INDEX]
    SD =  int(instance_data[0][START_DAY_INDEX])

    OPENING_TIME_WINDOW_ABBREV_KEY = 'O'
    CLOSING_TIME_WINDOW_ABBREV_KEY = 'C'
    TOTAL_TIME_KEY = 'Total Time'
    COLUMN_NAMES_ABBREV = ['i', 'x', 'y', 'd', 'S', 't',
                           'open_0', 'close_0', 'open_1', 'close_1',
                           'open_2', 'close_2', 'open_3', 'close_3',
                           'open_4', 'close_4', 'open_5', 'close_5',
                           'open_6', 'close_6', 'b']

    df = pd.DataFrame(instance_data[2:], columns=COLUMN_NAMES_ABBREV)

    df_ = df[['i', 'x', 'y', 'd', 'S', 't']+[c for c in df.columns if c[-1]==str(SD)]]
    columns = ['i', 'x', 'y', 'd', 'S', 't', OPENING_TIME_WINDOW_ABBREV_KEY, CLOSING_TIME_WINDOW_ABBREV_KEY]
    df_.columns=columns

    aux = pd.DataFrame([instance_data[1]], columns = ['i', 'x', 'y', 'd', 'S', 'O', 'C'])
    df = pd.concat([aux, df_], axis=0, sort=True).reset_index(drop=True)

    #add total time
    df[TOTAL_TIME_KEY] = 0
    df[TOTAL_TIME_KEY] = df.loc[0][CLOSING_TIME_WINDOW_ABBREV_KEY]

    return df
def read_instance_data(instance_name, path):

    """reads instance data"""
    PATH_TO_BENCHMARK_INSTANCES = path

    benchmark_file = '{path_to_benchmark_instances}/{instance}.txt' \
                     .format(path_to_benchmark_instances=PATH_TO_BENCHMARK_INSTANCES,
                             instance=instance_name)

    dfile = open(benchmark_file)
    data = [[float(x) for x in line.split()] for line in dfile]
    dfile.close()
    return data
def get_distance_matrix(instance_df, instance_type):
    """
    Distances between locations were rounded down to the first decimal
    for the Solomon instances and to the second decimal for the instances of Cordeau and Gavalas.
    """

    if instance_type in ['Solomon']:
        n_digits = 10.0

    elif instance_type in ['Cordeau', 'Gavalas']:
        n_digits = 100.0

    n = instance_df.shape[0]
    distm = np.zeros((n,n))
    x = instance_df.x.values
    y = instance_df.y.values

    for i in range(0, n-1):
        for j in range(i+1, n):
            distm[i,j] = np.floor(n_digits*(np.sqrt((x[i]-x[j])**2+(y[i]-y[j])**2)))/n_digits
            distm[j,i] = distm[i,j]

    return distm
def get_instance_data(instance, path, device):
    instance_type = get_instance_type(instance)
    df_inst = get_instance_df(instance, path, instance_type)
    distm = get_distance_matrix(df_inst, instance_type)
    raw_data = df_inst[['x', 'y', 'duration', 'ti', 'tf', 'prof', 'Total Time']].values
    raw_data = torch.FloatTensor(raw_data).to(device)
    raw_distm =  torch.FloatTensor(distm).to(device)

    return raw_data, raw_distm
def get_instance_df(instance_name, path, instance_type):

    """combine read instance, tests and parse to dataframe"""

    OPENING_TIME_WINDOW_ABBREV_KEY = 'O'
    CLOSING_TIME_WINDOW_ABBREV_KEY = 'C'
    TOTAL_TIME_KEY = 'Total Time'


    COLUMN_NAMES = ['vertex number', 'x coordinate', 'y coordinate',
    'service duration or visiting time', 'profit of the location',
    'not relevant 1', 'not relevant 2', 'not relevant 3',
    'opening of time window', 'closing of time window']
    COLUMN_NAMES = [s.replace(' ', '_') for s in COLUMN_NAMES]
    COLUMN_NAMES_ABBREV = ['i', 'x', 'y', 'd', 'S', 'f', 'a', 'list', 'O', 'C']
    VERTEX_NUMBER_COL = [i for i,n in enumerate(COLUMN_NAMES) if n=='vertex_number'][0]
    COLS_OF_INT = ['i', 'x', 'y', 'd', OPENING_TIME_WINDOW_ABBREV_KEY,
                   CLOSING_TIME_WINDOW_ABBREV_KEY, 'S', TOTAL_TIME_KEY]
    COLS_OF_INT_NEW_NAMES = ['i', 'x', 'y', 'duration', 'ti', 'tf', 'prof', TOTAL_TIME_KEY]

    standard2newnames_dict =  dict(((c, ca) for c, ca in zip(COLS_OF_INT, COLS_OF_INT_NEW_NAMES)))

    instance_data = read_instance_data(instance_name, path)

    # run tests
    test_n_vert_1(instance_data, instance_type)
    test_n_vert_2(instance_data, instance_type)
    # test if it's a single day (we are not considering TOPTW instances)
    test_n_vert_3(instance_data, instance_type)

    if instance_type=='Gavalas':
        df = parse_instance_data_Gavalas(instance_data)
    else:
        df = parse_instance_data(instance_data)

    #change column names
    COLS_OF_INT_NEW_NAMES = [standard2newnames_dict[s] for s in COLS_OF_INT]
    df_ = df[COLS_OF_INT].copy()
    df_.columns = COLS_OF_INT_NEW_NAMES
    df_['inst_name'] = instance_name
    df_['real_or_val'] = 'real'

    df_ = df_.append(df_.loc[0])
    return df_

def get_real_data(args, phase='train'):

    df = get_instance_df(args.instance, cf.BENCHMARK_INSTANCES_PATH,
                         args.instance_type)
    dist_mat = get_distance_matrix(df, args.instance_type)
    inp_real = df[['x', 'y', 'duration', 'ti', 'tf', 'prof', 'Total Time']].values

    if phase=='train':
        inp_real = [(torch.FloatTensor(inp_real).to(args.device),
                torch.tensor(inp_real[0, pcf.OPENING_TIME_WINDOW_IDX]).to(args.device),
                torch.FloatTensor(dist_mat).to(args.device))]

        new_inp_real = [(args.instance, inp_real[0])]
        return new_inp_real
    else:
        inp_real = [(torch.FloatTensor(inp_real).to(args.device),
                 torch.FloatTensor(dist_mat).to(args.device))]
        return inp_real
def get_val_data(args, phase='train'):
    path_string = os.path.abspath(os.getcwd())+'/{directory}/{file_name}'
    inp_val_path = os.path.normpath(path_string.format(directory=args.val_dir, file_name=args.val_set_pt_file))
    inp_val = torch.load(inp_val_path, map_location = args.map_location)

    new_inp_val = [(args.instance, inst_data) for inst_data in inp_val]

    if phase=='train':
        return new_inp_val
    else:

        return inp_val<jupyter_output><empty_output><jupyter_text>Feature Utils<jupyter_code>class DynamicFeatures():

    def __init__(self, args):
        super(DynamicFeatures, self).__init__()

        self.arrival_time_idx = pcf.ARRIVAL_TIME_IDX
        self.opening_time_window_idx = pcf.OPENING_TIME_WINDOW_IDX
        self.closing_time_window_idx = pcf.CLOSING_TIME_WINDOW_IDX
        self.device = args.device

    def make_dynamic_feat(self, data, current_time, current_poi_idx, dist_mat, batch_idx):

        num_dyn_feat = 8
        _ , sequence_size, input_size  = data.size()
        batch_size = batch_idx.shape[0]

        dyn_feat = torch.ones(batch_size, sequence_size, num_dyn_feat).to(self.device)

        tour_start_time = data[0, 0, self.opening_time_window_idx]
        max_tour_duration = data[0, 0, self.arrival_time_idx] - tour_start_time
        arrive_j_times = current_time + dist_mat[current_poi_idx]

        dyn_feat[:, :, 0] = (data[batch_idx, :, self.opening_time_window_idx] - current_time) / max_tour_duration
        dyn_feat[:, :, 1] = (data[batch_idx, :, self.closing_time_window_idx] - current_time) / max_tour_duration
        dyn_feat[:, :, 2] = (data[batch_idx, :, self.arrival_time_idx] - current_time) / max_tour_duration
        dyn_feat[:, :, 3] = (current_time - tour_start_time) / max_tour_duration


        dyn_feat[:, :, 4] = (arrive_j_times - tour_start_time) / max_tour_duration
        dyn_feat[:, :, 5] = (data[batch_idx, :, self.opening_time_window_idx] - arrive_j_times) / max_tour_duration
        dyn_feat[:, :, 6] = (data[batch_idx, :, self.closing_time_window_idx] - arrive_j_times) / max_tour_duration
        dyn_feat[:, :, 7] = (data[batch_idx, :, self.arrival_time_idx] - arrive_j_times) / max_tour_duration

        return dyn_feat<jupyter_output><empty_output><jupyter_text>Sampling Norm Utils<jupyter_code>def instance_dependent_norm_const(instance_raw_data):
    day_duration = int(instance_raw_data[:, pcf.CLOSING_TIME_WINDOW_IDX].max().item())
    t_max_real = int(instance_raw_data[0, pcf.ARRIVAL_TIME_IDX].item()) # max instance arrival time
    arrival_time_val_ub = t_max_real+int(pcf.SAMP_DAY_FRAC_INF*day_duration)
    Tmax = int(max(day_duration, arrival_time_val_ub)) # max possible time value
    Smax = int(torch.round(pcf.MULTIPLE_SCORE*instance_raw_data[1:-1, pcf.REWARD_IDX].max()).item()) # max score

    return Tmax, Smax
def sample_new_instance(inst_data, dist_mat, args):
    instance_type = args.instance_type
    sample_type  = args.sample_type

    if instance_type=='Solomon':
        n_digits = 10.0
        xy_inf = 0.
        xy_delta = 100.
    elif instance_type=='Cordeau':
        n_digits = 100.0
        xy_inf = -100.
        xy_delta = 200.
    elif instance_type=='Gavalas':
        n_digits = 100.0
        xy_inf = 0.
        xy_delta = 100.

    poit = inst_data.clone()
    n = inst_data.shape[0]

    prof = inst_data[1:-1, pcf.REWARD_IDX]
    durat_max = int(inst_data[1:-1, pcf.VIS_DURATION_TIME_IDX].max().item())

    day_duration = int(inst_data[:, pcf.CLOSING_TIME_WINDOW_IDX].max().item())

    t_init_real = int(inst_data[0, pcf.OPENING_TIME_WINDOW_IDX].item()) # starting time
    t_max_real = int(inst_data[0, pcf.ARRIVAL_TIME_IDX].item()) # max arrival time

    day_fract_inf = pcf.SAMP_DAY_FRAC_INF
    t_min = int(pcf.SAMP_DAY_FRAC_INF*day_duration)
    ub_t_init_val = pcf.UB_T_INIT_FRAC*day_duration
    lb_t_max_val = pcf.LB_T_MAX_FRAC*day_duration

    ub = int(torch.min(torch.tensor([ub_t_init_val,
                                     t_max_real+int(day_fract_inf*day_duration)])))
    t_init_val = torch.randint(int(t_init_real)-int(day_fract_inf*day_duration),
                               ub,
                               (1,))

    lb = int(torch.max(torch.tensor([lb_t_max_val, int(t_init_val)+t_min])))
    t_max_val = torch.randint(lb,
                              t_max_real+int(day_fract_inf*day_duration),
                              (1,))

    Smax = int(torch.round(pcf.MULTIPLE_SCORE*prof.max()).item())
    if sample_type == 'uni_samp':
        new_scores = torch.randint(1, Smax, (n,))
    elif sample_type == 'corr_samp':
        new_scores_unbound = (Smax/durat_max)*inst_data[:, pcf.VIS_DURATION_TIME_IDX] + pcf.CORR_SCORE_STD*torch.randn(n, device=args.device)
        new_scores = torch.round(torch.min(Smax*torch.ones(1, device=args.device),
                                           torch.max(torch.ones(n, device=args.device),
                                                     new_scores_unbound)))

    poit[:, pcf.REWARD_IDX] = new_scores

    #------------ correct first/last----------

    poit[0, pcf.REWARD_IDX] = 0 # starting point
    poit[n-1, pcf.REWARD_IDX] = 0 # ending point

    poit[0, pcf.X_COORDINATE_IDX] = float(xy_inf+xy_delta*torch.rand(1)) # starting point
    poit[n-1, pcf.X_COORDINATE_IDX] = poit[0, pcf.X_COORDINATE_IDX] # ending point

    poit[0, pcf.Y_COORDINATE_IDX] = float(xy_inf+xy_delta*torch.rand(1)) # starting point
    poit[n-1, pcf.Y_COORDINATE_IDX] = poit[0, pcf.Y_COORDINATE_IDX] # ending point

    poit[:, pcf.ARRIVAL_TIME_IDX] = t_max_val*torch.ones(n)
    poit[0, pcf.OPENING_TIME_WINDOW_IDX] = t_init_val*torch.ones(1)
    poit[n-1, pcf.OPENING_TIME_WINDOW_IDX] = t_init_val*torch.ones(1)
    poit[0, pcf.CLOSING_TIME_WINDOW_IDX] = t_max_val*torch.ones(1)
    poit[n-1, pcf.CLOSING_TIME_WINDOW_IDX] = t_max_val*torch.ones(1)

    start_time = t_init_val.clone().detach().to(args.device)
    dist_matn = dist_mat.clone()

    for j in range(1, n-1):
        dist_matn[0, j] = float(torch.floor(n_digits*(torch.sqrt((poit[0,0]-poit[j,0])**2+(poit[0,1]-poit[j,1])**2))).item()/n_digits)
        dist_matn[n-1, j] = dist_matn[0, j]

        dist_matn[j, 0] = dist_matn[0, j]
        dist_matn[j, n-1] = dist_matn[n-1, j]

    return poit, start_time, dist_matn

def data_scaler(data, norm_dic):
    datan = data.clone()
    datan[:, pcf.X_COORDINATE_IDX] /= pcf.X_MAX
    datan[:, pcf.Y_COORDINATE_IDX] /= pcf.X_MAX
    datan[:, pcf.VIS_DURATION_TIME_IDX] /= (datan[:, pcf.VIS_DURATION_TIME_IDX].max())
    datan[:, pcf.OPENING_TIME_WINDOW_IDX] /= norm_dic['Tmax']
    datan[:, pcf.CLOSING_TIME_WINDOW_IDX ] /= norm_dic['Tmax']
    datan[:, pcf.REWARD_IDX] /= norm_dic['Smax']
    datan[:, pcf.ARRIVAL_TIME_IDX] /= norm_dic['Tmax']

    return datan<jupyter_output><empty_output><jupyter_text>Train Utils<jupyter_code>def train_model(raw_data, raw_dist_mat, norm_dic, run_episode, opt, args):

    new_data, start_time, dist_mat = sample_new_instance(raw_data, raw_dist_mat, args)
    new_data_scaled = data_scaler(new_data, norm_dic[args.instance])
    bnew_data, bnew_data_scaled = samples2batch(new_data, new_data_scaled, args.batch_size)

    run_episode.train()
    opt.zero_grad()
    actions, log_prob, entropy, step_mask = run_episode(bnew_data, bnew_data_scaled, start_time, dist_mat, 'stochastic')

    rewards = reward_fn(new_data, actions, args.device)

    loss = 0

    av_rew = rewards.mean()
    min_rew = rewards.min()
    max_rew = rewards.max()

    advantage = (rewards - av_rew) #advantage

    res = advantage.unsqueeze(1)*log_prob + args.beta*entropy

    loss = -res[step_mask].sum()/args.batch_size

    loss.backward(retain_graph=False)
    torch.nn.utils.clip_grad_norm_(run_episode.neuralnet.parameters(), args.max_grad_norm)
    opt.step()

    return av_rew.item(), min_rew.item(), max_rew.item(), loss.item()

def samples2batch(new_data, new_data_scaled, batch_size):
    bnew_data = new_data.expand(batch_size, -1, -1)
    bnew_data_scaled = new_data_scaled.expand(batch_size, -1, -1)
    return bnew_data, bnew_data_scaled
def reward_fn(data, sample_solution, device):
    """
    Returns:
        Tensor of shape [batch_size] containing rewards
    """

    batch_size = sample_solution[0].shape[0]
    tour_reward = torch.zeros(batch_size, device=device)

    for act_id in sample_solution:
        tour_reward += data[act_id, pcf.REWARD_IDX].squeeze(0)

    return tour_reward

def exp_lr_scheduler(optimizer, epoch, init_lr=0.001, lr_decay_step=5000):
    """Decay learning rate by a factor of 0.96 every lr_decay_epoch epochs.
       Lower_bounded at 0.00001"""
    lr = init_lr * (0.96**(epoch // lr_decay_step))
    if lr < 0.00001:
        lr = 0.00001

    if epoch % lr_decay_step == 0:
        print('LR is set to {}'.format(lr))

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    return optimizer

def validation(inp_val, run_episode, inst_norm_dic, device):
    reward_val =  torch.tensor(0.0).to(device)
    rew_dict = {}
    for k, (inst_name, data) in enumerate(inp_val):
        inst_data, start_time, dist_mat = data
        rew = test_model(inst_data, start_time, dist_mat, inst_name, inst_norm_dic, run_episode, device)
        reward_val += rew
        key_str = inst_name + '_' + str(k)
        rew_dict[key_str] = rew

    return rew_dict, reward_val.item()/len(inp_val)

def test_model(data, start_time, dist_mat, inst, inst_norm_dic, run_episode, device):
    with torch.no_grad():
        data_scaled = data_scaler(data, inst_norm_dic[inst])
        bdata, bdata_scaled = data.unsqueeze(0), data_scaled.unsqueeze(0)
        actions, log_prob, entropy, step_mask = run_episode(bdata, bdata_scaled, start_time, dist_mat, 'greedy')
        reward = reward_fn(data, actions, device)

        return reward.item()<jupyter_output><empty_output><jupyter_text>Solution Construction<jupyter_code>class ModelUtils():
    def __init__(self, args):
        super(ModelUtils, self).__init__()

        self.device = args.device
        self.opening_time_window_idx = pcf.OPENING_TIME_WINDOW_IDX
        self.closing_time_window_idx = pcf.CLOSING_TIME_WINDOW_IDX
        self.vis_duration_time_idx = pcf.VIS_DURATION_TIME_IDX
        self.arrival_time_idx = pcf.ARRIVAL_TIME_IDX

    def feasibility_control(self, braw_inputs, mask, dist_mat, pres_act, present_time, batch_idx, first_step=False):

        done = False
        maskk = mask.clone()
        step_batch_size = batch_idx.shape[0]

        arrivej = dist_mat[pres_act] + present_time
        waitj = torch.max(torch.FloatTensor([0.0]).to(self.device), braw_inputs[:, :, self.opening_time_window_idx]-arrivej)

        c1 = arrivej + waitj <= braw_inputs[:, :, self.closing_time_window_idx]
        c2 = arrivej + waitj + braw_inputs[:, :, self.vis_duration_time_idx] + dist_mat[:, -1] <= braw_inputs[0, 0, self.arrival_time_idx]

        if not first_step:
            maskk[batch_idx, pres_act] = 0

        maskk[batch_idx] = maskk[batch_idx] * c1 * c2

        if maskk[:, -1].any() == 0:
            done = True
        return done, maskk


    def one_step_update(self, raw_inputs_b, dist_mat, pres_action, future_action, present_time, batch_idx, batch_size):

        present_time_b = torch.zeros(batch_size, 1, device=self.device)
        pres_actions_b = torch.zeros(batch_size, dtype=torch.int64, device=self.device)
        step_mask_b = torch.zeros(batch_size, 1, device=self.device, requires_grad=False, dtype=torch.bool)

        arrive_j = dist_mat[pres_action, future_action].unsqueeze(1) + present_time
        wait_j = torch.max(torch.FloatTensor([0.0]).to(self.device),
                           raw_inputs_b[batch_idx, future_action, self.opening_time_window_idx].unsqueeze(1)-arrive_j)
        present_time = arrive_j + wait_j + raw_inputs_b[batch_idx, future_action, self.vis_duration_time_idx].unsqueeze(1)

        present_time_b[batch_idx] = present_time

        pres_actions_b[batch_idx] = future_action
        step_mask_b[batch_idx] = 1

        return pres_actions_b, present_time_b, step_mask_b
class Lookahead():
    def __init__(self, args):
        super(Lookahead, self).__init__()

        self.device = args.device
        self.opening_time_window_idx = pcf.OPENING_TIME_WINDOW_IDX
        self.closing_time_window_idx = pcf.CLOSING_TIME_WINDOW_IDX
        self.vis_duration_time_idx = pcf.VIS_DURATION_TIME_IDX
        self.arrival_time_idx = pcf.ARRIVAL_TIME_IDX

    def adjacency_matrix(self, braw_inputs, mask, dist_mat, pres_act, present_time):
        # feasible neighborhood for each node
        maskk = mask.clone()
        step_batch_size, npoints = mask.shape

        #one step forward update
        arrivej = dist_mat[pres_act] + present_time
        farrivej = arrivej.view(step_batch_size, npoints)
        tw_start = braw_inputs[:, :, self.opening_time_window_idx]
        waitj = torch.max(torch.FloatTensor([0.0]).to(self.device), tw_start-farrivej)
        durat = braw_inputs[:, : , self.vis_duration_time_idx]

        fpresent_time = farrivej + waitj + durat
        fpres_act = torch.arange(0, npoints, device=self.device).expand(step_batch_size, -1)

        # feasible neighborhood for each node
        adj_mask = maskk.unsqueeze(1).repeat(1, npoints, 1)
        arrivej = dist_mat.expand(step_batch_size, -1, -1) + fpresent_time.unsqueeze(2)
        waitj = torch.max(torch.FloatTensor([0.0]).to(self.device), tw_start.unsqueeze(2)-arrivej)

        tw_end = braw_inputs[:, :, self.closing_time_window_idx]
        ttime = braw_inputs[:, 0, self.arrival_time_idx]

        dlast = dist_mat[:, -1].unsqueeze(0).expand(step_batch_size, -1)

        c1 = arrivej + waitj <= tw_end.unsqueeze(1)
        c2 = arrivej + waitj + durat.unsqueeze(1) + dlast.unsqueeze(1) <= ttime.unsqueeze(1).unsqueeze(1).expand(-1, npoints, npoints)
        adj_mask = adj_mask * c1 * c2

        # self-loop
        idx = torch.arange(0, npoints, device=self.device).expand(step_batch_size, -1)
        adj_mask[:, idx, idx] = 1

        return adj_mask
class RunEpisode(nn.Module):

    def __init__(self, neuralnet, args):
        super(RunEpisode, self).__init__()

        self.device = args.device
        self.neuralnet = neuralnet
        self.dyn_feat = DynamicFeatures(args)
        self.lookahead = Lookahead(args)
        self.mu = ModelUtils(args)

    def forward(self, binputs, bdata_scaled, start_time, dist_mat, infer_type):

        self.batch_size, sequence_size, input_size = binputs.size()

        h_0, c_0 = self.neuralnet.decoder.hidden_0

        dec_hidden = (h_0.expand(self.batch_size, -1), c_0.expand(self.batch_size, -1))

        mask = torch.ones(self.batch_size, sequence_size, device=self.device, requires_grad=False, dtype = torch.uint8)

        bpresent_time = start_time*torch.ones(self.batch_size, 1, device=self.device)

        llog_probs, lactions, lstep_mask, lentropy = [], [], [], []

        bpres_actions = torch.zeros(self.batch_size, dtype=torch.int64, device=self.device)

        batch_idx = torch.arange(0, self.batch_size, device=self.device)

        done, mask = self.mu.feasibility_control(binputs[batch_idx], mask, dist_mat, bpres_actions,
                                                 bpresent_time, batch_idx, first_step=True)

        adj_mask = self.lookahead.adjacency_matrix(binputs[batch_idx], mask, dist_mat, bpres_actions, bpresent_time)

        # encoder first forward pass
        bdyn_inputs = self.dyn_feat.make_dynamic_feat(binputs, bpresent_time, bpres_actions, dist_mat, batch_idx)
        emb1 = self.neuralnet.sta_emb(bdata_scaled)
        emb2 = self.neuralnet.dyn_emb(bdyn_inputs)
        enc_inputs = torch.cat((emb1,emb2), dim=2)

        _, _, enc_outputs = self.neuralnet(enc_inputs, enc_inputs, adj_mask, enc_inputs, dec_hidden, mask, first_step=True)

        decoder_input = enc_outputs[batch_idx, bpres_actions]

        done, mask = self.mu.feasibility_control(binputs[batch_idx], mask, dist_mat, bpres_actions, bpresent_time, batch_idx)
        adj_mask = self.lookahead.adjacency_matrix(binputs[batch_idx], mask, dist_mat, bpres_actions, bpresent_time)

        # encoder/decoder forward pass
        bdyn_inputs = self.dyn_feat.make_dynamic_feat(binputs, bpresent_time,
                                                      bpres_actions, dist_mat, batch_idx)
        emb2 = self.neuralnet.dyn_emb(bdyn_inputs)
        enc_inputs = torch.cat((emb1,emb2), dim=2)

        policy, dec_hidden, enc_outputs = self.neuralnet(enc_inputs, enc_outputs, adj_mask, decoder_input, dec_hidden, mask)

        lactions.append(bpres_actions)

        # Starting the trip
        while not done:

            future_actions, log_probs, entropy = self.select_actions(policy, infer_type)

            bpres_actions, bpresent_time, bstep_mask = self.mu.one_step_update(binputs, dist_mat, bpres_actions[batch_idx],
                                                                               future_actions, bpresent_time[batch_idx],
                                                                               batch_idx, self.batch_size)

            blog_probs = torch.zeros(self.batch_size, 1, dtype=torch.float32).to(self.device)
            blog_probs[batch_idx] = log_probs.unsqueeze(1)

            bentropy = torch.zeros(self.batch_size,1,dtype=torch.float32).to(self.device)
            bentropy[batch_idx] = entropy.unsqueeze(1)

            llog_probs.append(blog_probs)
            lactions.append(bpres_actions)
            lstep_mask.append(bstep_mask)
            lentropy.append(bentropy)

            done, mask = self.mu.feasibility_control(binputs[batch_idx], mask, dist_mat,
                                                     bpres_actions[batch_idx], bpresent_time[batch_idx],
                                                     batch_idx)

            if done: break
            sub_batch_idx = torch.nonzero(mask[batch_idx][:,-1], as_tuple=False).squeeze(1)

            batch_idx = torch.nonzero(mask[:,-1], as_tuple=False).squeeze(1)

            adj_mask = self.lookahead.adjacency_matrix(binputs[batch_idx], mask[batch_idx], dist_mat, bpres_actions[batch_idx], bpresent_time[batch_idx])

            #update decoder input and hidden
            decoder_input = enc_outputs[sub_batch_idx, bpres_actions[sub_batch_idx]]
            dec_hidden = (dec_hidden[0][sub_batch_idx], dec_hidden[1][sub_batch_idx])

            # encoder/decoder forward pass
            bdyn_inputs = self.dyn_feat.make_dynamic_feat(binputs, bpresent_time[batch_idx], bpres_actions[batch_idx], dist_mat, batch_idx)
            emb2 = self.neuralnet.dyn_emb(bdyn_inputs)
            enc_inputs = torch.cat((emb1[batch_idx],emb2), dim=2)

            policy, dec_hidden, enc_outputs = self.neuralnet(enc_inputs, enc_outputs[sub_batch_idx], adj_mask, decoder_input, dec_hidden, mask[batch_idx])

        return lactions, torch.cat(llog_probs, dim=1), torch.cat(lentropy, dim=1), torch.cat(lstep_mask, dim=1)


    def select_actions(self, policy, infer_type):

        if infer_type == 'stochastic':
            m = Categorical(policy)
            act_ind = m.sample()
            log_select =  m.log_prob(act_ind)
            poli_entro = m.entropy()
        elif infer_type == 'greedy':
            prob, act_ind = torch.max(policy, 1)
            log_select =  prob.log()
            poli_entro =  torch.zeros(self.batch_size, requires_grad=False).to(self.device)

        return act_ind, log_select, poli_entro<jupyter_output><empty_output><jupyter_text>Run<jupyter_code>def train_loop(inp_val, inp_real, raw_data, run_episode, args):

    raw_data, raw_dist_mat = inp_real[0][1][0], inp_real[0][1][2]
    reward_total, min_reward_total, max_reward_total, loss_total = 0, 0, 0, 0
    training_history = []
    model_opt = optim.Adam(run_episode.neuralnet.parameters(), lr=args.lr)
    step_dict = {}

    for epoch in tqdm(range(1,args.nepocs+1)):

        avg_reward, min_reward, max_reward, loss = train_model(raw_data,
                                                                  raw_dist_mat,
                                                                  norm_dic,
                                                                  run_episode,
                                                                  model_opt,
                                                                  args)

        reward_total += avg_reward
        min_reward_total += min_reward
        max_reward_total += max_reward
        loss_total += loss

        exp_lr_scheduler(model_opt, epoch, init_lr=args.lr)


        if epoch==0 or epoch % args.nprint == 0:
            print("Epoch %s" % str(epoch))
            rew_dict, avg_reward_val = validation(inp_val, run_episode, norm_dic, args.device)
            _, avg_reward_real  = validation(inp_real, run_episode, norm_dic, args.device)
            step_dict[epoch] = rew_dict

            if epoch == 0:
                avg_loss = loss_total
                avg_reward_total = reward_total
                avg_min_reward_total = min_reward_total
                avg_max_reward_total = max_reward_total

                training_history.append([epoch, reward_total, min_reward_total, max_reward_total,
                                         avg_reward_val, avg_reward_real, loss_total])

            else:
                avg_loss = loss_total / args.nprint
                avg_reward_total = reward_total / args.nprint
                avg_min_reward_total = min_reward_total / args.nprint
                avg_max_reward_total = max_reward_total / args.nprint


                training_history.append([epoch, avg_reward_total, avg_min_reward_total, avg_max_reward_total,
                                         avg_reward_val, avg_reward_real, avg_loss])


            print(N_DASHES*'-')
            print("Average total loss: %s" % avg_loss)
            print("Average train mean reward: %s" % avg_reward_total)
            print("Average train max reward: %s" % avg_max_reward_total)
            print("Average train min reward: %s" % avg_min_reward_total)
            print("Validation reward: %2.3f"  % (avg_reward_val))
            print("Real instance reward: %2.3f"  % (avg_reward_real))

            reward_total, min_reward_total, max_reward_total, loss_total = 0, 0, 0, 0

        if epoch % args.nsave == 0 and not args.debug:
            print('saving model')
            torch.save(run_episode.neuralnet.state_dict(), args.save_w_dir+'/model_'+str(epoch)+'.pkl')

    return training_history

args = DotMap(dict(
    batch_size=32, 
    beta=0.01, 
    debug=False, 
    device='cpu',
    device_name='cpu',
    ff_dim=256, 
    instance='t101', 
    instance_type='Gavalas', 
    lr=0.0001, 
    map_location={'cpu': 'cpu'}, 
    max_grad_norm=1,
    model_name='testing_1', 
    n_heads=8, 
    n_layers=2, 
    ndfeatures=8, 
    nepocs=1, 
    nfeatures=7, 
    nprint=1, 
    nsave=1, 
    rnn_hidden=128, 
    sample_type='corr_samp', 
    save_h_dir='./results/t101/outputs/model_testing_1_corr_samp', 
    save_h_file='training_history.csv', 
    save_w_dir='./results/t101/model_w/model_testing_1_corr_samp', 
    seed=2925, 
    use_checkpoint=False, 
    val_dir='./data/generated//t101', 
    val_set_pt_file='inp_val_corr_samp.pt'
))
# for reproducibility"
N_DASHES = 40
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if str(args.device) in ['cuda', 'cuda:0', 'cuda:1']:
    torch.cuda.manual_seed(args.seed)

# create directories for saving
if not args.debug:
    os.makedirs(args.save_h_dir) if not os.path.exists(args.save_h_dir) else None
    os.makedirs(args.save_w_dir) if not os.path.exists(args.save_w_dir) else None

# get val and real instance scores
inp_real = get_real_data(args, phase='train')
inp_val = get_val_data(args, phase='train')

# get Tmax and Smax
raw_data = inp_real[0][1][0]
Tmax, Smax = instance_dependent_norm_const(raw_data)
norm_dic = {args.instance: {'Tmax': Tmax, 'Smax': Smax}}
def save_training_history(training_history, file_path, args):
    TRAIN_HIST_COLUMNS = ['epoch', 'avg_reward_train',
        'min_reward_train', 'max_reward_train', 'avg_reward_val_'+args.sample_type,
        'avg_reward_real', 'tloss_train']

    df = pd.DataFrame(training_history, columns = TRAIN_HIST_COLUMNS)
    df.to_csv(file_path, index=False)
    return

# train
model = RecPointerNetwork(args.nfeatures, args.ndfeatures, args.rnn_hidden, args).to(args.device)
run_episode = RunEpisode(model, args)

training_history = train_loop(inp_val, inp_real, norm_dic, run_episode, args)

# save
if not args.debug:
    file_path = '%s/%s' % (args.save_h_dir, args.save_h_file)
    save_training_history(training_history, file_path, args)<jupyter_output>0%|                                                                                            | 0/1 [00:00<?, ?it/s]
{'fraction_non_alphanumeric': 0.11568919461976691, 'fraction_numerical': 0.011488546911566555, 'mean_word_length': 3.056381244606386, 'pattern_counts': {'":': 0, '<': 38, '<?xml version=': 0, '>': 33, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 40, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '252468', 'license': 'MIT', 'n_tokens_mistral': 15654, 'n_tokens_neox': 15094, 'n_words': 2791, 'path': '.ipynb_checkpoints/Learn_OPTW_Train-checkpoint.ipynb', 'repo_name': 'arsalanakhter/optw_rl', 'size': '42303'}