Upload models/apm.py with huggingface_hub

models/apm.py

@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class APM(nn.Module):
+    """Action Proposal Module: predicts the action bridging two latent states."""
+
+    def __init__(self, z_dim: int = 4, action_dim: int = 4, hidden_dim: int = 256):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(2 * z_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, action_dim),
+        )
+        self.z_dim = z_dim
+        self.action_dim = action_dim
+
+    def forward(self, z_i: torch.Tensor, z_j: torch.Tensor) -> torch.Tensor:
+        return self.net(torch.cat([z_i, z_j], dim=-1))
+
+
+def apm_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return F.mse_loss(pred, target)