| 372 | |
| 373 | # add context token to the encoder, put the context token in the decoder's inputs |
| 374 | def context_token_forward( |
| 375 | self, |
| 376 | states, |
| 377 | actions, |
| 378 | goals, |
| 379 | sigma |
| 380 | ): |
| 381 | |
| 382 | if len(states.size()) != 3: |
| 383 | states = states.unsqueeze(0) |
| 384 | |
| 385 | # t for the states does not mean the time, but the number of inputs tokens |
| 386 | b, t, dim = states.size() |
| 387 | _, t_a, _ = actions.size() |
| 388 | |
| 389 | if self.goal_conditioned: |
| 390 | goal_embed = self.goal_emb(goals) |
| 391 | goal_x = self.drop(goal_embed + self.pos_emb[:, :self.goal_seq_len, :]) |
| 392 | |
| 393 | state_embed = self.tok_emb(states) |
| 394 | state_x = self.drop(state_embed + self.pos_emb[:, self.goal_seq_len:(self.goal_seq_len + t), :]) |
| 395 | |
| 396 | context_token = self.context_embed.weight.unsqueeze(0).repeat(b, 1, 1) |
| 397 | |
| 398 | action_embed = self.action_emb(actions) |
| 399 | action_x = self.drop(action_embed + self.pos_emb[:, (self.goal_seq_len + t):(self.goal_seq_len + t + t_a), :]) |
| 400 | |
| 401 | emb_t = self.sigma_emb(sigma) |
| 402 | |
| 403 | if self.goal_conditioned: |
| 404 | input_seq = torch.cat([goal_x, state_x, context_token], dim=1) |
| 405 | else: |
| 406 | input_seq = torch.cat([state_x, context_token], dim=1) |
| 407 | |
| 408 | # adaLN conditioning |
| 409 | if self.use_ada_conditioning: |
| 410 | encoder_output = self.encoder(input_seq)[:, -1:, :] |
| 411 | emb_t = emb_t + encoder_output |
| 412 | decoder_output = self.decoder(action_x, emb_t) |
| 413 | else: |
| 414 | input_seq = torch.cat([emb_t, input_seq], dim=1) |
| 415 | encoder_output = self.encoder(input_seq)[:, -1:, :] |
| 416 | decoder_output = self.decoder(action_x, encoder_output) |
| 417 | |
| 418 | pred_actions = self.action_pred(decoder_output[:, -self.action_seq_len:, :]) |
| 419 | |
| 420 | return pred_actions |
| 421 | |
| 422 | def forward( |
| 423 | self, |