File size: 5,521 Bytes
be5548b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions.categorical import Categorical
import torch_ac

 
from utils.other import init_params




class RandomTalkingMultiHeadedACModel(nn.Module, torch_ac.RecurrentACModel):
    def __init__(self, obs_space, action_space, use_memory=False, use_text=False, use_dialogue=False):
        super().__init__()

        # Decide which components are enabled
        self.use_text = use_text
        self.use_dialogue = use_dialogue
        self.use_memory = use_memory

        # multi dim
        if action_space.shape == ():
            raise ValueError("The action space is not multi modal. Use ACModel instead.")

        self.n_primitive_actions = action_space.nvec[0] + 1
        self.talk_action = int(self.n_primitive_actions) - 1
        self.n_utterance_actions = action_space.nvec[1:]
        self.env_action_space = action_space
        self.model_raw_action_space = spaces.MultiDiscrete([self.n_primitive_actions, *self.n_utterance_actions])

        # Define image embedding
        self.image_conv = nn.Sequential(
            nn.Conv2d(3, 16, (2, 2)),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),
            nn.Conv2d(16, 32, (2, 2)),
            nn.ReLU(),
            nn.Conv2d(32, 64, (2, 2)),
            nn.ReLU()
        )
        n = obs_space["image"][0]
        m = obs_space["image"][1]
        self.image_embedding_size = ((n-1)//2-2)*((m-1)//2-2)*64

        # Define memory
        if self.use_memory:
            self.memory_rnn = nn.LSTMCell(self.image_embedding_size, self.semi_memory_size)

        if self.use_text or self.use_dialogue:
            self.word_embedding_size = 32
            self.word_embedding = nn.Embedding(obs_space["text"], self.word_embedding_size)

        # Define text embedding
        if self.use_text:
            self.text_embedding_size = 128
            self.text_rnn = nn.GRU(self.word_embedding_size, self.text_embedding_size, batch_first=True)

        # Define dialogue embedding
        if self.use_dialogue:
            self.dialogue_embedding_size = 128
            self.dialogue_rnn = nn.GRU(self.word_embedding_size, self.dialogue_embedding_size, batch_first=True)

        # Resize image embedding
        self.embedding_size = self.semi_memory_size

        if self.use_text:
            self.embedding_size += self.text_embedding_size

        if self.use_dialogue:
            self.embedding_size += self.dialogue_embedding_size

        # Define actor's model
        self.actor = nn.Sequential(
            nn.Linear(self.embedding_size, 64),
            nn.Tanh(),
            nn.Linear(64, self.n_primitive_actions)
        )

        # Define critic's model
        self.critic = nn.Sequential(
            nn.Linear(self.embedding_size, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )


        # Initialize parameters correctly
        self.apply(init_params)

    @property
    def memory_size(self):
        return 2*self.semi_memory_size

    @property
    def semi_memory_size(self):
        return self.image_embedding_size

    def forward(self, obs, memory):
        x = obs.image.transpose(1, 3).transpose(2, 3)
        x = self.image_conv(x)

        batch_size = x.shape[0]
        x = x.reshape(batch_size, -1)

        if self.use_memory:
            hidden = (memory[:, :self.semi_memory_size], memory[:, self.semi_memory_size:])
            hidden = self.memory_rnn(x, hidden)
            embedding = hidden[0]
            memory = torch.cat(hidden, dim=1)
        else:
            embedding = x

        if self.use_text:
            embed_text = self._get_embed_text(obs.text)
            embedding = torch.cat((embedding, embed_text), dim=1)

        if self.use_dialogue:
            embed_dial = self._get_embed_dialogue(obs.dialogue)
            embedding = torch.cat((embedding, embed_dial), dim=1)

        x = self.actor(embedding)
        primtive_actions_dist = Categorical(logits=F.log_softmax(x, dim=1))

        x = self.critic(embedding)
        value = x.squeeze(1)

        # construct utterance action distributions, for this model they are radndom
        utterance_actions_dists = [Categorical(logits=torch.ones((batch_size, n), requires_grad=False)) for n in self.n_utterance_actions]

        dist = [primtive_actions_dist] + utterance_actions_dists

        return dist, value, memory

    def sample_action(self, dist):
        return torch.stack([d.sample() for d in dist], dim=1)

    def calculate_log_probs(self, dist, action):
        return torch.stack([d.log_prob(action[:, i]) for i, d in enumerate(dist)], dim=1)

    def calculate_action_masks(self, action):
        talk_mask = action[:, 0] == self.talk_action
        mask = torch.stack(
            (torch.ones_like(talk_mask), talk_mask, talk_mask),
            dim=1).detach()

        assert action.shape == mask.shape

        return mask

    def construct_final_action(self, action):
        act_mask = action[:, 0] != self.n_primitive_actions - 1

        nan_mask = np.array([
            np.array([1, np.nan, np.nan]) if t else np.array([np.nan, 1, 1]) for t in act_mask
        ])

        action = nan_mask*action

        return action

    def _get_embed_text(self, text):
        _, hidden = self.text_rnn(self.word_embedding(text))
        return hidden[-1]

    def _get_embed_dialogue(self, dial):
        _, hidden = self.dialogue_rnn(self.word_embedding(dial))
        return hidden[-1]