Added more models. Fixed number of actions

Cubevoid · Apr 18, 2024 · b500fb6 · b500fb6
1 parent 0523043
commit b500fb6
Show file tree

Hide file tree

Showing 10 changed files with 2 additions and 2 deletions.
diff --git a/models/trained/Freeway/1713396710_Predictor.pth b/models/trained/Freeway/1713396710_Predictor.pth
diff --git a/models/trained/Freeway/1713396710_feat_extract.pth b/models/trained/Freeway/1713396710_feat_extract.pth
diff --git a/models/trained/Freeway/1713396773_PredictorBaseline.pth b/models/trained/Freeway/1713396773_PredictorBaseline.pth
diff --git a/models/trained/Freeway/1713396773_feat_extract.pth b/models/trained/Freeway/1713396773_feat_extract.pth
diff --git a/models/trained/Freeway/1713396990_ResidualPredictor.pth b/models/trained/Freeway/1713396990_ResidualPredictor.pth
diff --git a/models/trained/Freeway/1713396990_feat_extract.pth b/models/trained/Freeway/1713396990_feat_extract.pth
diff --git a/models/trained/Pong/1713396916_ResidualPredictor.pth b/models/trained/Pong/1713396916_ResidualPredictor.pth
diff --git a/models/trained/Pong/1713396916_feat_extract.pth b/models/trained/Pong/1713396916_feat_extract.pth
diff --git a/src/model/predictor.py b/src/model/predictor.py
@@ -17,7 +17,7 @@ def __init__(self, input_size: int = 128, hidden_size: int = 32, output_size: in
         self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers)
         self.time_mlp = nn.Sequential(nn.Linear(output_size, output_size))
         self.pred_mlp = nn.Sequential(nn.Linear(output_size, output_size), nn.ReLU(), nn.Linear(output_size, 2))
-        self.action_embedding = nn.Embedding(16, embed_dim)
+        self.action_embedding = nn.Embedding(18, embed_dim)
         self.embedding = nn.Sequential(nn.Linear(output_size+embed_dim, output_size), nn.ReLU())
 
     def forward(self, x: torch.Tensor, curr_pos: torch.Tensor, actions: torch.Tensor) -> torch.Tensor:  # pylint: disable = unused-argument

diff --git a/src/model/predictor_baseline.py b/src/model/predictor_baseline.py
@@ -8,7 +8,7 @@ def __init__(self, input_size: int = 128, time_steps: int = 5, embed_dim: int =
         self.encoder = nn.Sequential(nn.Linear(input_size, input_size), nn.ReLU(), nn.Linear(input_size, input_size))
         self.next_state = nn.Sequential(nn.Linear(input_size, input_size), nn.ReLU(), nn.Linear(input_size, input_size))
         self.output = nn.Sequential(nn.Linear(input_size, input_size), nn.ReLU(), nn.Linear(input_size, 2))
-        self.action_embedding = nn.Embedding(16, embed_dim)
+        self.action_embedding = nn.Embedding(18, embed_dim)
         self.embedding = nn.Sequential(nn.Linear(input_size+embed_dim, input_size), nn.ReLU())
 
     def forward(self, x: torch.Tensor, curr_pos: torch.Tensor, actions: torch.Tensor) -> torch.Tensor: