Upload 4 files

ant/pwm_torch_seperate/encoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ae346f29cece02fefdc82310919c046f5c930a3eb165a5f521dfccbd1388154
+size 407381

ant/pwm_torch_seperate/random_torch.py

@@ -0,0 +1,113 @@
+"""
+generate_random_wm.py
+
+Creates random (but valid) world-model networks:
+- Encoder
+- Transition model
+- Reward model
+
+Saves class-compatible state_dict weights.
+"""
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# -------------------------
+# Config (Ant-style)
+# -------------------------
+OBS_DIM = 105
+ACTION_DIM = 8
+LATENT_DIM = 32
+HIDDEN_DIM = 256
+SEED = 42
+
+OUT_DIR = "weights"
+
+
+# -------------------------
+# Models
+# -------------------------
+class Encoder(nn.Module):
+    def __init__(self, obs_dim: int, latent_dim: int, hidden_dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(obs_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, latent_dim),
+        )
+
+    def forward(self, obs):
+        return self.net(obs)
+
+
+class TransitionModel(nn.Module):
+    def __init__(self, latent_dim: int, action_dim: int, hidden_dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(latent_dim + action_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, latent_dim),
+        )
+
+    def forward(self, z, action):
+        x = torch.cat([z, action], dim=-1)
+        return self.net(x)
+
+
+class RewardModel(nn.Module):
+    def __init__(self, latent_dim: int, action_dim: int, hidden_dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(latent_dim + action_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, 1),
+        )
+
+    def forward(self, z, action):
+        x = torch.cat([z, action], dim=-1)
+        return self.net(x).squeeze(-1)
+
+
+# -------------------------
+# Initialization
+# -------------------------
+def init_weights(m):
+    if isinstance(m, nn.Linear):
+        nn.init.orthogonal_(m.weight)
+        nn.init.zeros_(m.bias)
+
+
+# -------------------------
+# Main
+# -------------------------
+def main():
+    torch.manual_seed(SEED)
+
+    encoder = Encoder(OBS_DIM, LATENT_DIM, HIDDEN_DIM)
+    transition = TransitionModel(LATENT_DIM, ACTION_DIM, HIDDEN_DIM)
+    reward = RewardModel(LATENT_DIM, ACTION_DIM, HIDDEN_DIM)
+
+    encoder.apply(init_weights)
+    transition.apply(init_weights)
+    reward.apply(init_weights)
+
+    os.makedirs(OUT_DIR, exist_ok=True)
+
+    torch.save(encoder.state_dict(), f"{OUT_DIR}/encoder.pth")
+    torch.save(transition.state_dict(), f"{OUT_DIR}/transition.pth")
+    torch.save(reward.state_dict(), f"{OUT_DIR}/reward.pth")
+
+    print("✅ Random world-model weights saved:")
+    print(f"  {OUT_DIR}/encoder.pth")
+    print(f"  {OUT_DIR}/transition.pth")
+    print(f"  {OUT_DIR}/reward.pth")
+
+
+if __name__ == "__main__":
+    main()

ant/pwm_torch_seperate/reward.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad796e55dfaebcba46a0baaec334c8e996794aace6950df447e6168f735426ab
+size 45331

ant/pwm_torch_seperate/transition.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd055e79288d6f6d3e10d83dab550a1dfb843261d44b7af029b6c7b5d4d4e5cd
+size 341177