| 16 | # Non-diffusion based encoder-decoder model |
| 17 | class EncDec(nn.Module): |
| 18 | def __init__( |
| 19 | self, |
| 20 | encoder: DictConfig, |
| 21 | decoder: DictConfig, |
| 22 | state_dim: int, |
| 23 | goal_dim: int, |
| 24 | action_dim: int, |
| 25 | device: str, |
| 26 | goal_conditioned: bool, |
| 27 | embed_dim: int, |
| 28 | embed_pdrob: float, |
| 29 | goal_seq_len: int, |
| 30 | obs_seq_len: int, |
| 31 | action_seq_len: int, |
| 32 | linear_output: bool = False, |
| 33 | forward_type: str = 'cross_attn' # cross_attn, context_token |
| 34 | ): |
| 35 | super().__init__() |
| 36 | |
| 37 | self.encoder = hydra.utils.instantiate(encoder) |
| 38 | self.decoder = hydra.utils.instantiate(decoder) |
| 39 | |
| 40 | self.device = device |
| 41 | |
| 42 | # mainly used for language condition or goal image condition |
| 43 | self.goal_conditioned = goal_conditioned |
| 44 | if not goal_conditioned: |
| 45 | goal_seq_len = 0 |
| 46 | |
| 47 | # the seq_size is the number of tokens in the input sequence |
| 48 | self.seq_size = goal_seq_len + obs_seq_len + action_seq_len |
| 49 | |
| 50 | # linear embedding for the state |
| 51 | self.tok_emb = nn.Linear(state_dim, embed_dim) |
| 52 | |
| 53 | # linear embedding for the goal |
| 54 | self.goal_emb = nn.Linear(goal_dim, embed_dim) |
| 55 | |
| 56 | # position embedding |
| 57 | self.pos_emb = nn.Parameter(torch.zeros(1, self.seq_size, embed_dim)) |
| 58 | self.drop = nn.Dropout(embed_pdrob) |
| 59 | self.drop.to(self.device) |
| 60 | |
| 61 | # get an action embedding |
| 62 | self.query_embed = nn.Embedding(action_seq_len, embed_dim) |
| 63 | |
| 64 | self.action_dim = action_dim |
| 65 | self.obs_dim = state_dim |
| 66 | self.embed_dim = embed_dim |
| 67 | |
| 68 | self.goal_seq_len = goal_seq_len |
| 69 | self.obs_seq_len = obs_seq_len |
| 70 | self.action_seq_len = action_seq_len |
| 71 | |
| 72 | self.forward_type = forward_type |
| 73 | |
| 74 | if self.forward_type != 'cross_attn': |
| 75 | self.context_embed = nn.Embedding(1, embed_dim) |