add trained weights

Author: kywch

pufferlib/ocean/g2048/g2048.c

@@ -1,6 +1,25 @@
 #include "g2048.h"
 #include "puffernet.h"
 
+// Network with hidden size 256. Should go to puffernet
+LinearLSTM* make_linearlstm_256(Weights* weights, int num_agents, int input_dim, int logit_sizes[], int num_actions) {
+    LinearLSTM* net = calloc(1, sizeof(LinearLSTM));
+    net->num_agents = num_agents;
+    net->obs = calloc(num_agents*input_dim, sizeof(float));
+    int hidden_dim = 256;
+    net->encoder = make_linear(weights, num_agents, input_dim, hidden_dim);
+    net->gelu1 = make_gelu(num_agents, hidden_dim);
+    int atn_sum = 0;
+    for (int i = 0; i < num_actions; i++) {
+        atn_sum += logit_sizes[i];
+    }
+    net->actor = make_linear(weights, num_agents, hidden_dim, atn_sum);
+    net->value_fn = make_linear(weights, num_agents, hidden_dim, 1);
+    net->lstm = make_lstm(weights, num_agents, hidden_dim, hidden_dim);
+    net->multidiscrete = make_multidiscrete(num_agents, logit_sizes, num_actions);
+    return net;
+}
+
 int main() {
     srand(time(NULL));
     Game env;
@@ -14,9 +33,9 @@ int main() {
     env.actions = actions;
     env.rewards = rewards;
 
-    Weights* weights = load_weights("resources/g2048/g2048_weights.bin", 134917);
+    Weights* weights = load_weights("resources/g2048/g2048_weights.bin", 531973);
     int logit_sizes[1] = {4};
-    LinearLSTM* net = make_linearlstm(weights, 1, 16, logit_sizes, 1);
+    LinearLSTM* net = make_linearlstm_256(weights, 1, 16, logit_sizes, 1);
     c_reset(&env);
     c_render(&env);
 

pufferlib/ocean/torch.py

@@ -190,47 +190,6 @@ def decode_actions(self, hidden):
         return action, value
 
 
-class G2048(nn.Module):
-    def __init__(self, env, cnn_channels=32, hidden_size=128):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.is_continuous = False
-
-        self.cnn = nn.Sequential(
-            pufferlib.pytorch.layer_init(
-                nn.Conv2d(1, cnn_channels, 2, stride=1)),
-            nn.GELU(),
-            pufferlib.pytorch.layer_init(
-                nn.Conv2d(cnn_channels, cnn_channels, 2, stride=1)),
-            nn.Flatten(),
-            nn.GELU(),
-            pufferlib.pytorch.layer_init(
-            nn.Linear(128, hidden_size), std=0.01),
-        )
-
-        self.decoder = pufferlib.pytorch.layer_init(
-            nn.Linear(hidden_size, env.single_action_space.n), std=0.01)
-        self.value = pufferlib.pytorch.layer_init(
-            nn.Linear(hidden_size, 1), std=1)
-
-    def forward_eval(self, observations, state=None):
-        hidden = self.encode_observations(observations)
-        actions, value = self.decode_actions(hidden)
-        return actions, value
-
-    def forward(self, x, state=None):
-        return self.forward_eval(x, state)
-
-    def encode_observations(self, observations, state=None):
-        #observations = F.one_hot(observations.long(), 16).view(-1, 16, 4, 4).float()
-        observations = observations.float().view(-1, 1, 4, 4)
-        return self.cnn(observations)
-
-    def decode_actions(self, hidden):
-        action = self.decoder(hidden)
-        value = self.value(hidden)
-        return action, value
-
 class Snake(nn.Module):
     def __init__(self, env, cnn_channels=32, hidden_size=128):
         super().__init__()