policy_model.py

import numpy as np

import torch
import torch.nn as nn

from models.modules import build_mlp, SoftAttention, PositionalEncoding, create_mask, proj_masking


class Regretful(nn.Module):
    """
    The model for the regretful agent (CVPR 2019)

    The Regretful Agent: Heuristic-Aided Navigation through Progress Estimation
    GitHub: https://github.com/chihyaoma/regretful-agent
    Project: https://chihyaoma.github.io/project/2019/02/25/regretful.html
    """

    def __init__(self, opts, img_fc_dim, img_fc_use_batchnorm, img_dropout, img_feat_input_dim,
                 rnn_hidden_size, rnn_dropout, max_len, fc_bias=True, max_navigable=16):
        super(Regretful, self).__init__()

        self.opts = opts
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.max_navigable = max_navigable

        proj_navigable_kwargs = {
            'input_dim': img_feat_input_dim,
            'hidden_dims': img_fc_dim,
            'use_batchnorm': img_fc_use_batchnorm,
            'dropout': img_dropout,
            'fc_bias': fc_bias
        }
        self.proj_navigable_mlp = build_mlp(**proj_navigable_kwargs)

        self.h0_fc = nn.Linear(rnn_hidden_size, img_fc_dim[-1])

        self.h1_fc = nn.Linear(rnn_hidden_size, rnn_hidden_size, bias=fc_bias)
        self.positional_encoding = PositionalEncoding(rnn_hidden_size, dropout=0.1, max_len=max_len)
        self.soft_attn = SoftAttention()

        self.dropout = nn.Dropout(p=rnn_dropout)
        self.lstm = nn.LSTMCell(img_fc_dim[-1] * 2 + rnn_hidden_size, rnn_hidden_size)

        self.logit_fc = nn.Linear(rnn_hidden_size * 2, img_fc_dim[-1])

        self.h2_fc_lstm = nn.Linear(rnn_hidden_size + img_fc_dim[-1], rnn_hidden_size, bias=fc_bias)

        self.critic_fc = nn.Linear(max_len + rnn_hidden_size, 1)
        self.tanh = nn.Tanh()
        self.sigmoid = nn.Sigmoid()

        self.critic_valueDiff_fc = nn.Linear(1, 2)
        self.relu = nn.ReLU(inplace=True)
        self.softmax = nn.Softmax(dim=1)

        self.move_fc = nn.Linear(img_fc_dim[-1], img_fc_dim[-1] + opts.tiled_len)

        self.num_predefined_action = 1

    def block_oscillation(self, batch_size, navigable_mask, oscillation_index, block_oscillation_index):

        navigable_mask[torch.LongTensor(range(batch_size)), np.array(oscillation_index) + 1] = \
        navigable_mask[torch.LongTensor(range(batch_size)), np.array(oscillation_index) + 1] * (1 - torch.Tensor(block_oscillation_index).to(self.device))

        assert (navigable_mask.sum(1) > 0).all(), "All actions are blocked ...? "

        return navigable_mask

    def forward(self, img_feat, navigable_feat, pre_feat, pre_value, h_0, c_0, ctx, navigable_index=None, navigable_idx_to_previous=None,
                oscillation_index=None, block_oscillation_index=None, ctx_mask=None, seq_lengths=None,
                prevent_oscillation=False, prevent_rollback=False, is_training=None):
        """
        forward passing the network of regretful agent

        :param img_feat: (batch_size, 36, d-dim feat)
        :param navigable_feat: (batch_size, max number of navigable direction, d-dim)
        :param pre_feat: (batch_size, d-dim)
        :param pre_value: (batch_size, 1)
        :param h_0: (batch_size, d-dim)
        :param c_0: (batch_size, d-dim)
        :param ctx: (batch_size, max instruction length, d-dim)
        :param navigable_index: list of list, index for navigable directions for each sample in the mini-batch
        :param navigable_idx_to_previous: list
        :param oscillation_index: list
        :param block_oscillation_index: list
        :param ctx_mask: (batch_size, max instruction length)
        :param seq_lengths: list
        :param prevent_oscillation: 1 or 0
        :param prevent_rollback: 1 or 0
        :param is_training: True or False
        """

        batch_size, num_imgs, feat_dim = img_feat.size()

        # creating a mask to block out non-navigable directions, due to batch processing
        index_length = [len(_index) + self.num_predefined_action for _index in navigable_index]
        navigable_mask = create_mask(batch_size, self.max_navigable, index_length)

        # prevent rollback action as a sanity check. See Table 3 in the paper.
        if prevent_rollback and not is_training:
            navigable_mask[torch.LongTensor(range(batch_size)), np.array(oscillation_index) + 1] = \
                navigable_mask[torch.LongTensor(range(batch_size)), np.array(oscillation_index) + 1] * (
                            1 - (torch.Tensor(index_length) > 2)).float().to(self.device)

        # block the navigable direction that leads to oscillation
        if 1 in block_oscillation_index and prevent_oscillation:
            navigable_mask = self.block_oscillation(batch_size, navigable_mask, oscillation_index,
                                                    block_oscillation_index)

        # get navigable features without attached markers for visual grounding
        navigable_feat_no_visited = navigable_feat[:, :, :-self.opts.tiled_len]
        proj_navigable_feat = proj_masking(navigable_feat_no_visited, self.proj_navigable_mlp, navigable_mask)
        proj_pre_feat = self.proj_navigable_mlp(pre_feat[:, :-self.opts.tiled_len])
        weighted_img_feat, img_attn = self.soft_attn(self.h0_fc(h_0), proj_navigable_feat, mask=navigable_mask)

        # positional encoding instruction embeddings and textual grounding
        positioned_ctx = self.positional_encoding(ctx)
        weighted_ctx, ctx_attn = self.soft_attn(self.h1_fc(h_0), positioned_ctx, mask=ctx_mask)

        # merge info into one LSTM to be carry through time
        concat_input = torch.cat((proj_pre_feat, weighted_img_feat, weighted_ctx), 1)
        h_1, c_1 = self.lstm(concat_input, (h_0, c_0))
        h_1_drop = self.dropout(h_1)

        # =========== forward and rollback embeddings ===========
        m_forward = self.logit_fc(torch.cat((weighted_ctx, h_1_drop), dim=1))
        m_rollback = proj_navigable_feat[torch.LongTensor(range(batch_size)),
                     np.array(navigable_idx_to_previous) + 1, :]

        # =========== Progress Monitor ===========
        concat_value_input = self.h2_fc_lstm(torch.cat((h_0, weighted_img_feat), 1))
        h_1_value = self.dropout(torch.sigmoid(concat_value_input) * torch.tanh(c_1))
        critics_input = torch.cat((ctx_attn, h_1_value), dim=1)

        if self.opts.monitor_sigmoid:
            value = self.sigmoid(self.critic_fc(critics_input))
        else:
            value = self.tanh(self.critic_fc(critics_input))

        # =========== Progress Marker ===========
        value_detached = value.detach()
        value_for_marker = value_detached.unsqueeze(1).repeat(1, self.max_navigable, self.opts.tiled_len) * \
                           navigable_mask.unsqueeze(2).repeat(1, 1, self.opts.tiled_len)

        navigable_visited_feat = value_for_marker - navigable_feat[:, :, -self.opts.tiled_len:]

        proj_navigable_feat_visited = torch.cat((proj_navigable_feat, navigable_visited_feat), dim=2)

        # =========== Regret Module ===========
        rollback_forward = torch.cat((m_rollback.unsqueeze(1), m_forward.unsqueeze(1)), dim=1)
        rollback_forward_logit = self.critic_valueDiff_fc(value_detached - pre_value)
        rollback_forward_attn = self.softmax(rollback_forward_logit)
        m_forward_rollback = torch.bmm(rollback_forward_attn.unsqueeze(1), rollback_forward).squeeze(1)

        # =========== Action selection with Progress Marker ===========
        logit = torch.bmm(proj_navigable_feat_visited, self.move_fc(m_forward_rollback).unsqueeze(2)).squeeze(2)


        return h_1, c_1, img_attn, ctx_attn, rollback_forward_attn, logit, rollback_forward_logit, value, navigable_mask


class SelfMonitoring(nn.Module):
    """
    The model for the self-monitoring agent (ICLR 2019)

    Self-Monitoring Navigation Agent via Auxiliary Progress Estimation
    arXiv: https://arxiv.org/abs/1901.03035
    GitHub: https://github.com/chihyaoma/selfmonitoring-agent
    Project: https://chihyaoma.github.io/project/2018/09/27/selfmonitoring.html

    """

    def __init__(self, opts, img_fc_dim, img_fc_use_batchnorm, img_dropout, img_feat_input_dim,
                 rnn_hidden_size, rnn_dropout, max_len, fc_bias=True, max_navigable=16):
        super(SelfMonitoring, self).__init__()

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.max_navigable = max_navigable

        proj_navigable_kwargs = {
            'input_dim': img_feat_input_dim,
            'hidden_dims': img_fc_dim,
            'use_batchnorm': img_fc_use_batchnorm,
            'dropout': img_dropout,
            'fc_bias': fc_bias,
            'relu': opts.mlp_relu
        }
        self.proj_navigable_mlp = build_mlp(**proj_navigable_kwargs)

        self.h0_fc = nn.Linear(rnn_hidden_size, img_fc_dim[-1], bias=fc_bias)
        self.h1_fc = nn.Linear(rnn_hidden_size, rnn_hidden_size, bias=fc_bias)

        self.soft_attn = SoftAttention()

        self.dropout = nn.Dropout(p=rnn_dropout)

        self.lstm = nn.LSTMCell(img_fc_dim[-1] * 2 + rnn_hidden_size, rnn_hidden_size)

        self.lang_position = PositionalEncoding(rnn_hidden_size, dropout=0.1, max_len=max_len)

        self.logit_fc = nn.Linear(rnn_hidden_size * 2, img_fc_dim[-1])
        self.h2_fc_lstm = nn.Linear(rnn_hidden_size + img_fc_dim[-1], rnn_hidden_size, bias=fc_bias)

        if opts.monitor_sigmoid:
            self.critic = nn.Sequential(
                nn.Linear(max_len + rnn_hidden_size, 1),
                nn.Sigmoid()
            )
        else:
            self.critic = nn.Sequential(
                nn.Linear(max_len + rnn_hidden_size, 1),
                nn.Tanh()
            )

        self.num_predefined_action = 1

    def forward(self, img_feat, navigable_feat, pre_feat, question, h_0, c_0, ctx, pre_ctx_attend,
                navigable_index=None, ctx_mask=None):
        """ Takes a single step in the decoder

        img_feat: batch x 36 x feature_size
        navigable_feat: batch x max_navigable x feature_size
        pre_feat: previous attended feature, batch x feature_size
        question: this should be a single vector representing instruction
        ctx: batch x seq_len x dim
        navigable_index: list of list
        ctx_mask: batch x seq_len - indices to be masked
        """
        batch_size, num_imgs, feat_dim = img_feat.size()

        index_length = [len(_index) + self.num_predefined_action for _index in navigable_index]
        navigable_mask = create_mask(batch_size, self.max_navigable, index_length)

        proj_navigable_feat = proj_masking(navigable_feat, self.proj_navigable_mlp, navigable_mask)
        proj_pre_feat = self.proj_navigable_mlp(pre_feat)
        positioned_ctx = self.lang_position(ctx)

        weighted_ctx, ctx_attn = self.soft_attn(self.h1_fc(h_0), positioned_ctx, mask=ctx_mask)

        weighted_img_feat, img_attn = self.soft_attn(self.h0_fc(h_0), proj_navigable_feat, mask=navigable_mask)

        # merge info into one LSTM to be carry through time
        concat_input = torch.cat((proj_pre_feat, weighted_img_feat, weighted_ctx), 1)

        h_1, c_1 = self.lstm(concat_input, (h_0, c_0))
        h_1_drop = self.dropout(h_1)

        # policy network
        h_tilde = self.logit_fc(torch.cat((weighted_ctx, h_1_drop), dim=1))
        logit = torch.bmm(proj_navigable_feat, h_tilde.unsqueeze(2)).squeeze(2)

        # value estimation
        concat_value_input = self.h2_fc_lstm(torch.cat((h_0, weighted_img_feat), 1))

        h_1_value = self.dropout(torch.sigmoid(concat_value_input) * torch.tanh(c_1))

        value = self.critic(torch.cat((ctx_attn, h_1_value), dim=1))

        return h_1, c_1, weighted_ctx, img_attn, ctx_attn, logit, value, navigable_mask

class SpeakerFollowerBaseline(nn.Module):
    """
    The baseline model for the speaker-follower (NeurIPS 2018).
    Note that this implementation is only for the basic speaker-follow without Pragmatic Inference

    arXiv: https://arxiv.org/abs/1806.02724
    """

    def __init__(self, opts, img_fc_dim, img_fc_use_batchnorm, img_dropout, img_feat_input_dim,
                 rnn_hidden_size, rnn_dropout, max_len, fc_bias=True, max_navigable=16):
        super(SpeakerFollowerBaseline, self).__init__()

        self.max_navigable = max_navigable

        self.proj_img_mlp = nn.Linear(img_feat_input_dim, img_fc_dim[-1], bias=fc_bias)

        self.proj_navigable_mlp = nn.Linear(img_feat_input_dim, img_fc_dim[-1], bias=fc_bias)

        self.h0_fc = nn.Linear(rnn_hidden_size, img_fc_dim[-1], bias=False)

        self.soft_attn = SoftAttention()

        self.dropout = nn.Dropout(p=rnn_dropout)

        self.lstm = nn.LSTMCell(img_feat_input_dim * 2, rnn_hidden_size)

        self.h1_fc = nn.Linear(rnn_hidden_size, rnn_hidden_size, bias=False)

        self.proj_out = nn.Linear(rnn_hidden_size, img_fc_dim[-1], bias=fc_bias)

    def forward(self, img_feat, navigable_feat, pre_feat, h_0, c_0, ctx, navigable_index=None, ctx_mask=None):
        """ Takes a single step in the decoder LSTM.

        img_feat: batch x 36 x feature_size
        navigable_feat: batch x max_navigable x feature_size
        h_0: batch x hidden_size
        c_0: batch x hidden_size
        ctx: batch x seq_len x dim
        navigable_index: list of list
        ctx_mask: batch x seq_len - indices to be masked
        """
        batch_size, num_imgs, feat_dim = img_feat.size()

        # add 1 because the navigable index yet count in "stay" location
        # but navigable feature does include the "stay" location at [:,0,:]
        index_length = [len(_index)+1 for _index in navigable_index]
        navigable_mask = create_mask(batch_size, self.max_navigable, index_length)

        proj_img_feat = proj_masking(img_feat, self.proj_img_mlp)

        proj_navigable_feat = proj_masking(navigable_feat, self.proj_navigable_mlp, navigable_mask)

        weighted_img_feat, _ = self.soft_attn(self.h0_fc(h_0), proj_img_feat, img_feat)

        concat_input = torch.cat((pre_feat, weighted_img_feat), 1)

        h_1, c_1 = self.lstm(self.dropout(concat_input), (h_0, c_0))

        h_1_drop = self.dropout(h_1)

        # use attention on language instruction
        weighted_context, ctx_attn = self.soft_attn(self.h1_fc(h_1_drop), self.dropout(ctx), mask=ctx_mask)
        h_tilde = self.proj_out(weighted_context)

        logit = torch.bmm(proj_navigable_feat, h_tilde.unsqueeze(2)).squeeze(2)

        return h_1, c_1, ctx_attn, logit, navigable_mask