actions: shape (B, T, action_dim) timesteps: shape (B,) -- a single scalar per batch item returns: shape (B, T, hidden_size)
(self, actions, timesteps)
| 67 | self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) |
| 68 | |
| 69 | def forward(self, actions, timesteps): |
| 70 | """ |
| 71 | actions: shape (B, T, action_dim) |
| 72 | timesteps: shape (B,) -- a single scalar per batch item |
| 73 | returns: shape (B, T, hidden_size) |
| 74 | """ |
| 75 | B, T, _ = actions.shape |
| 76 | |
| 77 | # 1) Expand each batch's single scalar time 'tau' across all T steps |
| 78 | # so that shape => (B, T) |
| 79 | # e.g. if timesteps is (B,), replicate across T |
| 80 | if timesteps.dim() == 1 and timesteps.shape[0] == B: |
| 81 | # shape (B,) => (B,T) |
| 82 | timesteps = timesteps.unsqueeze(1).expand(-1, T) |
| 83 | else: |
| 84 | raise ValueError( |
| 85 | "Expected `timesteps` to have shape (B,) so we can replicate across T." |
| 86 | ) |
| 87 | |
| 88 | # 2) Standard action MLP step for shape => (B, T, w) |
| 89 | a_emb = self.W1(actions) |
| 90 | |
| 91 | # 3) Get the sinusoidal encoding (B, T, w) |
| 92 | tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) |
| 93 | |
| 94 | # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish |
| 95 | x = torch.cat([a_emb, tau_emb], dim=-1) |
| 96 | x = swish(self.W2(x)) |
| 97 | |
| 98 | # 5) Finally W3 => (B, T, w) |
| 99 | x = self.W3(x) |
| 100 | |
| 101 | return x |