Produces a sinusoidal encoding of shape (B, T, w) given timesteps of shape (B, T).
| 22 | |
| 23 | |
| 24 | class SinusoidalPositionalEncoding(nn.Module): |
| 25 | """ |
| 26 | Produces a sinusoidal encoding of shape (B, T, w) |
| 27 | given timesteps of shape (B, T). |
| 28 | """ |
| 29 | |
| 30 | def __init__(self, embedding_dim): |
| 31 | super().__init__() |
| 32 | self.embedding_dim = embedding_dim |
| 33 | |
| 34 | def forward(self, timesteps): |
| 35 | # timesteps: shape (B, T) |
| 36 | # We'll compute sin/cos frequencies across dim T |
| 37 | timesteps = timesteps.float() # ensure float |
| 38 | |
| 39 | B, T = timesteps.shape |
| 40 | device = timesteps.device |
| 41 | |
| 42 | half_dim = self.embedding_dim // 2 |
| 43 | # typical log space frequencies for sinusoidal encoding |
| 44 | exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * ( |
| 45 | torch.log(torch.tensor(10000.0)) / half_dim |
| 46 | ) |
| 47 | # Expand timesteps to (B, T, 1) then multiply |
| 48 | freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) |
| 49 | |
| 50 | sin = torch.sin(freqs) |
| 51 | cos = torch.cos(freqs) |
| 52 | enc = torch.cat([sin, cos], dim=-1) # (B, T, w) |
| 53 | |
| 54 | return enc |
| 55 | |
| 56 | |
| 57 | class ActionEncoder(nn.Module): |