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