add test script

rlcube/rlcube/envs/__init__.py

@@ -1,18 +1,3 @@
 from .cube2 import Cube2Env
 
 __all__ = ["Cube2Env"]
-
-if __name__ == "__main__":
-    print("Testing Cube2Env.from_obs")
-    env = Cube2Env()
-    env2 = Cube2Env.from_obs(env.obs())
-    print(env2.state, env2.is_solved())
-    env2.print_js_code()
-    print()
-
-    print("Testing Cube2Env.adjacent_obs")
-    env = Cube2Env()
-    adjacent_obs = env.adjacent_obs()
-    for i in range(12):
-        env = Cube2Env.from_obs(adjacent_obs[i])
-        env.print_js_code()

rlcube/rlcube/envs/cube2.py

@@ -242,7 +242,7 @@ class Cube2Env(gym.Env):
         for i in range(6):
             for j in range(4):
                 label = int(self.state[i, j])
-                zeros = np.zeros(6)
+                zeros = np.zeros(6, dtype=np.int8)
                 zeros[label] = 1
                 one_hots.append(zeros)
         return np.array(one_hots)

rlcube/rlcube/envs/test.py

@@ -0,0 +1,17 @@
+from rlcube.envs.cube2 import Cube2Env
+
+
+if __name__ == "__main__":
+    print("Testing Cube2Env.from_obs")
+    env = Cube2Env()
+    env2 = Cube2Env.from_obs(env.obs())
+    print(env2.state, env2.is_solved())
+    env2.print_js_code()
+    print()
+
+    print("Testing Cube2Env.adjacent_obs")
+    env = Cube2Env()
+    adjacent_obs = env.adjacent_obs()
+    for i in range(12):
+        env = Cube2Env.from_obs(adjacent_obs[i])
+        env.print_js_code()

rlcube/rlcube/models/dataset.py

@@ -1,5 +1,5 @@
 from torch.utils.data import Dataset
-from rlcube.envs.cube2 import Cube2
+from rlcube.envs import Cube2Env
 import numpy as np
 import torch
 from tqdm import tqdm
@@ -12,13 +12,13 @@ def create_dataset(
     neighbors = []
     D = []
     for _ in tqdm(range(num_envs)):
-        env = Cube2()
+        env = Cube2Env()
         obs, _ = env.reset()
         for _ in range(num_steps):
             action = env.action_space.sample()
             obs, _, _, _, _ = env.step(action)
             states.append(obs)
-            neighbors.append(env.neighbors())
+            neighbors.append(env.adjacent_obs())
             D.append(env.step_count)
     states = np.array(states)
     neighbors = np.array(neighbors)

rlcube/rlcube/models/models.py

@@ -2,7 +2,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch
 from tensordict import TensorDict
-from rlcube.envs.cube2 import Cube2
+from rlcube.envs.cube2 import Cube2Env
 import numpy as np
 
 
@@ -20,8 +20,8 @@ class Reward(nn.Module):
         solved = face_solved.all(dim=1)
         return torch.where(
             solved,
-            torch.tensor(1, device=batch_obs.device, dtype=batch_obs.dtype),
-            torch.tensor(-1, device=batch_obs.device, dtype=batch_obs.dtype),
+            torch.tensor(1, device=batch_obs.device, dtype=torch.float32),
+            torch.tensor(-1, device=batch_obs.device, dtype=torch.float32),
         )
 
 
@@ -82,7 +82,7 @@ class DNN(nn.Module):
 
 if __name__ == "__main__":
     print("Testing RewardNet")
-    env = Cube2()
+    env = Cube2Env()
     obs, _ = env.reset()
     obs1, _, _, _, _ = env.step(1)
     obs2, _, _, _, _ = env.step(2)

rlcube/rlcube/models/test.py

@@ -0,0 +1,17 @@
+from rlcube.envs.cube2 import Cube2Env
+from rlcube.models.models import Reward
+import torch
+
+if __name__ == "__main__":
+    print("Testing Reward")
+    env = Cube2Env()
+    adjacent_obs = torch.tensor(env.adjacent_obs())
+    reward = Reward()(adjacent_obs)
+    print("Reward:", reward)
+    obs, _, _, _, _ = env.step(0)
+    env1 = Cube2Env.from_obs(obs)
+    adjacent_obs = env1.adjacent_obs()
+    for i in range(12):
+        env2 = Cube2Env.from_obs(adjacent_obs[i])
+        env2.print_js_code()
+        print(env2.is_solved())