| 55 | |
| 56 | |
| 57 | class ActionEncoder(nn.Module): |
| 58 | def __init__(self, action_dim, hidden_size): |
| 59 | super().__init__() |
| 60 | self.hidden_size = hidden_size |
| 61 | |
| 62 | # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} |
| 63 | self.W1 = nn.Linear(action_dim, hidden_size) # (d -> w) |
| 64 | self.W2 = nn.Linear(2 * hidden_size, hidden_size) # (2w -> w) |
| 65 | self.W3 = nn.Linear(hidden_size, hidden_size) # (w -> w) |
| 66 | |
| 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 |
nothing calls this directly
no outgoing calls
no test coverage detected