chore: import upstream snapshot with attribution

2026-07-13 12:36:51 +08:00
commit d731240295
27 changed files with 3078 additions and 0 deletions
@@ -0,0 +1,50 @@
+# Copyright 2022 Twitter, Inc and Zhendong Wang.
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from agents.helpers import SinusoidalPosEmb
+
+
+class MLP(nn.Module):
+    """
+    MLP Model
+    """
+    def __init__(self,
+                 state_dim,
+                 action_dim,
+                 device,
+                 t_dim=16):
+
+        super(MLP, self).__init__()
+        self.device = device
+
+        self.time_mlp = nn.Sequential(
+            SinusoidalPosEmb(t_dim),
+            nn.Linear(t_dim, t_dim * 2),
+            nn.Mish(),
+            nn.Linear(t_dim * 2, t_dim),
+        )
+
+        input_dim = state_dim + action_dim + t_dim
+        self.mid_layer = nn.Sequential(nn.Linear(input_dim, 256),
+                                       nn.Mish(),
+                                       nn.Linear(256, 256),
+                                       nn.Mish(),
+                                       nn.Linear(256, 256),
+                                       nn.Mish())
+
+        self.final_layer = nn.Linear(256, action_dim)
+
+    def forward(self, x, time, state):
+
+        t = self.time_mlp(time)
+        x = torch.cat([x, t, state], dim=1)
+        x = self.mid_layer(x)
+
+        return self.final_layer(x)
+
+