(
self,
encoder: DictConfig,
decoder: DictConfig,
state_dim: int,
goal_dim: int,
action_dim: int,
device: str,
goal_conditioned: bool,
embed_dim: int,
embed_pdrob: float,
goal_seq_len: int,
obs_seq_len: int,
action_seq_len: int,
linear_output: bool = False,
use_ada_conditioning: bool = False,
diffusion_type: str = "beso", # ddpm, beso or rf
forward_type: str = 'cross_attn' # cross_attn, context_token
)
| 200 | # Diffusion based decoder-only model, we need time embedding and noisy antions inputs here |
| 201 | class Noise_EncDec(nn.Module): |
| 202 | def __init__( |
| 203 | self, |
| 204 | encoder: DictConfig, |
| 205 | decoder: DictConfig, |
| 206 | state_dim: int, |
| 207 | goal_dim: int, |
| 208 | action_dim: int, |
| 209 | device: str, |
| 210 | goal_conditioned: bool, |
| 211 | embed_dim: int, |
| 212 | embed_pdrob: float, |
| 213 | goal_seq_len: int, |
| 214 | obs_seq_len: int, |
| 215 | action_seq_len: int, |
| 216 | linear_output: bool = False, |
| 217 | use_ada_conditioning: bool = False, |
| 218 | diffusion_type: str = "beso", # ddpm, beso or rf |
| 219 | forward_type: str = 'cross_attn' # cross_attn, context_token |
| 220 | ): |
| 221 | super().__init__() |
| 222 | |
| 223 | self.encoder = hydra.utils.instantiate(encoder) |
| 224 | self.decoder = hydra.utils.instantiate(decoder) |
| 225 | |
| 226 | self.device = device |
| 227 | |
| 228 | # mainly used for language condition or goal image condition |
| 229 | self.goal_conditioned = goal_conditioned |
| 230 | if not goal_conditioned: |
| 231 | goal_seq_len = 0 |
| 232 | |
| 233 | # the seq_size is the number of tokens in the input sequence |
| 234 | self.seq_size = goal_seq_len + obs_seq_len + action_seq_len |
| 235 | |
| 236 | # linear embedding for the state |
| 237 | self.tok_emb = nn.Linear(state_dim, embed_dim) |
| 238 | |
| 239 | # linear embedding for the goal |
| 240 | self.goal_emb = nn.Linear(goal_dim, embed_dim) |
| 241 | |
| 242 | # linear embedding for the action |
| 243 | self.action_emb = nn.Linear(action_dim, embed_dim) |
| 244 | |
| 245 | self.diffusion_type = diffusion_type |
| 246 | |
| 247 | if diffusion_type == "beso": |
| 248 | self.sigma_emb = BESO_TimeEmbedding(embed_dim) |
| 249 | elif diffusion_type == "rf": |
| 250 | self.sigma_emb = RF_TimeEmbedding(embed_dim) |
| 251 | else: |
| 252 | raise ValueError(f"Diffusion type {diffusion_type} is not supported") |
| 253 | |
| 254 | # position embedding |
| 255 | self.pos_emb = nn.Parameter(torch.zeros(1, self.seq_size, embed_dim)) |
| 256 | self.drop = nn.Dropout(embed_pdrob) |
| 257 | self.drop.to(self.device) |
| 258 | |
| 259 | self.action_dim = action_dim |
nothing calls this directly
no test coverage detected