| 151 | |
| 152 | # add context token to the encoder, put the context token in the decoder's inputs |
| 153 | def context_token_forward(self, states, goals): |
| 154 | |
| 155 | if len(states.size()) != 3: |
| 156 | states = states.unsqueeze(0) |
| 157 | |
| 158 | b, t, dim = states.size() |
| 159 | |
| 160 | if self.goal_conditioned: |
| 161 | goal_embed = self.goal_emb(goals) |
| 162 | goal_x = self.drop(goal_embed + self.pos_emb[:, :self.goal_seq_len, :]) |
| 163 | |
| 164 | state_embed = self.tok_emb(states) |
| 165 | state_x = self.drop(state_embed + self.pos_emb[:, self.goal_seq_len:(self.goal_seq_len + t), :]) |
| 166 | |
| 167 | context_token = self.context_embed.weight.unsqueeze(0).repeat(b, 1, 1) |
| 168 | |
| 169 | if self.goal_conditioned: |
| 170 | input_seq = torch.cat([goal_x, state_x, context_token], dim=1) |
| 171 | else: |
| 172 | input_seq = torch.cat([state_x, context_token], dim=1) |
| 173 | |
| 174 | # only output the context token |
| 175 | encoder_output = self.encoder(input_seq)[:, -1:, :] |
| 176 | |
| 177 | # decode the action sequence with cross attention over the encoder output |
| 178 | action_seq = self.query_embed.weight.unsqueeze(0).repeat(b, 1, 1) |
| 179 | |
| 180 | decoder_output = self.decoder(action_seq, encoder_output) |
| 181 | |
| 182 | pred_actions = self.action_pred(decoder_output[:, -self.action_seq_len:, :]) |
| 183 | |
| 184 | return pred_actions |
| 185 | |
| 186 | def forward( |
| 187 | self, |