implement transition

rlcube/rlcube/envs/cube2.py

@@ -6,8 +6,8 @@ F = 0
 B = 1
 R = 2
 L = 3
-T = 4
-B = 5
+U = 4
+D = 5
 
 class Cube2(gym.Env):
     def __init__(self):
@@ -24,13 +24,13 @@ class Cube2(gym.Env):
         self.state[1] = np.ones(4) * B
         self.state[2] = np.ones(4) * R
         self.state[3] = np.ones(4) * L 
-        self.state[4] = np.ones(4) * T
-        self.state[5] = np.ones(4) * B
-        # shuffle_steps =self.np_random.integers(0, 20)
-        # for i in range(shuffle_steps):
-        #     self.step(self.action_space.sample())
+        self.state[4] = np.ones(4) * U
+        self.state[5] = np.ones(4) * D
+        shuffle_steps =self.np_random.integers(0, 20)
+        for i in range(shuffle_steps):
+            self.step(self.action_space.sample())
         self.step_count = 0
-        return self.state, {}
+        return self._get_obs(), {}
     
     def step(self, action):
         self.step_count += 1
@@ -42,114 +42,114 @@ class Cube2(gym.Env):
             new_state[F, 1] = self.state[F, 0]
             new_state[F, 2] = self.state[F, 3]
             new_state[F, 3] = self.state[F, 1]
-            new_state[R, 1] = self.state[T, 3]
-            new_state[R, 3] = self.state[T, 1]
-            new_state[L, 1] = self.state[B, 3]
-            new_state[L, 3] = self.state[B, 1]
-            new_state[T, 1] = self.state[L, 1]
-            new_state[T, 3] = self.state[L, 3]
-            new_state[B, 1] = self.state[R, 1]
-            new_state[B, 3] = self.state[R, 3]
+            new_state[R, 1] = self.state[U, 3]
+            new_state[R, 3] = self.state[U, 1]
+            new_state[L, 1] = self.state[D, 3]
+            new_state[L, 3] = self.state[D, 1]
+            new_state[U, 1] = self.state[L, 1]
+            new_state[U, 3] = self.state[L, 3]
+            new_state[D, 1] = self.state[R, 1]
+            new_state[D, 3] = self.state[R, 3]
         # Front Counter-Clockwise
-        if action == 1:
+        elif action == 1:
             new_state[F, 0] = self.state[F, 1]
             new_state[F, 1] = self.state[F, 3]
             new_state[F, 2] = self.state[F, 0]
             new_state[F, 3] = self.state[F, 2]
-            new_state[R, 1] = self.state[B, 1]
-            new_state[R, 3] = self.state[B, 3]
-            new_state[L, 1] = self.state[T, 1]
-            new_state[L, 3] = self.state[T, 3]
-            new_state[T, 1] = self.state[R, 3]
-            new_state[T, 3] = self.state[R, 1]
-            new_state[B, 1] = self.state[L, 3]
-            new_state[B, 3] = self.state[L, 1]
+            new_state[R, 1] = self.state[D, 1]
+            new_state[R, 3] = self.state[D, 3]
+            new_state[L, 1] = self.state[U, 1]
+            new_state[L, 3] = self.state[U, 3]
+            new_state[U, 1] = self.state[R, 3]
+            new_state[U, 3] = self.state[R, 1]
+            new_state[D, 1] = self.state[L, 3]
+            new_state[D, 3] = self.state[L, 1]
         # Back Clockwise
-        if action == 2:
+        elif action == 2:
             new_state[B, 0] = self.state[B, 1]
             new_state[B, 1] = self.state[B, 3]
             new_state[B, 2] = self.state[B, 0]
             new_state[B, 3] = self.state[B, 2]
-            new_state[R, 0] = self.state[B, 0]
-            new_state[R, 2] = self.state[B, 2]
-            new_state[L, 0] = self.state[T, 0]
-            new_state[L, 2] = self.state[T, 2]
-            new_state[T, 0] = self.state[R, 2]
-            new_state[T, 2] = self.state[R, 0]
-            new_state[B, 0] = self.state[L, 2]
-            new_state[B, 2] = self.state[L, 0]
+            new_state[R, 0] = self.state[D, 0]
+            new_state[R, 2] = self.state[D, 2]
+            new_state[L, 0] = self.state[U, 0]
+            new_state[L, 2] = self.state[U, 2]
+            new_state[U, 0] = self.state[R, 2]
+            new_state[U, 2] = self.state[R, 0]
+            new_state[D, 0] = self.state[L, 2]
+            new_state[D, 2] = self.state[L, 0]
         # Back Counter-Clockwise
-        if action == 3:
+        elif action == 3:
             new_state[B, 0] = self.state[B, 2]
             new_state[B, 1] = self.state[B, 0]
             new_state[B, 2] = self.state[B, 3]
             new_state[B, 3] = self.state[B, 1]
-            new_state[R, 0] = self.state[T, 2]
-            new_state[R, 2] = self.state[T, 0]
-            new_state[L, 0] = self.state[B, 2]
-            new_state[L, 2] = self.state[B, 0]
-            new_state[T, 0] = self.state[L, 0]
-            new_state[T, 2] = self.state[L, 2]
-            new_state[B, 0] = self.state[R, 0]
-            new_state[B, 2] = self.state[R, 2]
+            new_state[R, 0] = self.state[U, 2]
+            new_state[R, 2] = self.state[U, 0]
+            new_state[L, 0] = self.state[D, 2]
+            new_state[L, 2] = self.state[D, 0]
+            new_state[U, 0] = self.state[L, 0]
+            new_state[U, 2] = self.state[L, 2]
+            new_state[D, 0] = self.state[R, 0]
+            new_state[D, 2] = self.state[R, 2]
         # Right Clockwise
-        if action == 4:
-            new_state[F, 2] = self.state[B, 2]
-            new_state[F, 3] = self.state[B, 3]
-            new_state[B, 2] = self.state[T, 2]
-            new_state[B, 3] = self.state[T, 3]
+        elif action == 4:
+            new_state[F, 2] = self.state[D, 2]
+            new_state[F, 3] = self.state[D, 3]
+            new_state[B, 2] = self.state[U, 2]
+            new_state[B, 3] = self.state[U, 3]
             new_state[R, 0] = self.state[R, 2]
             new_state[R, 1] = self.state[R, 0]
             new_state[R, 2] = self.state[R, 3]
             new_state[R, 3] = self.state[R, 1]
-            new_state[T, 2] = self.state[F, 3]
-            new_state[T, 3] = self.state[F, 2]
-            new_state[B, 2] = self.state[B, 3]
-            new_state[B, 3] = self.state[B, 2]
+            new_state[U, 2] = self.state[F, 3]
+            new_state[U, 3] = self.state[F, 2]
+            new_state[D, 2] = self.state[B, 3]
+            new_state[D, 3] = self.state[B, 2]
         # Right Counter-Clockwise
-        if action == 5:
-            new_state[F, 2] = self.state[T, 3]
-            new_state[F, 3] = self.state[T, 2]
-            new_state[B, 2] = self.state[B, 3]
-            new_state[B, 3] = self.state[B, 2]
+        elif action == 5:
+            new_state[F, 2] = self.state[U, 3]
+            new_state[F, 3] = self.state[U, 2]
+            new_state[B, 2] = self.state[D, 3]
+            new_state[B, 3] = self.state[D, 2]
             new_state[R, 0] = self.state[R, 1]
             new_state[R, 1] = self.state[R, 3]
             new_state[R, 2] = self.state[R, 0]
             new_state[R, 3] = self.state[R, 2]
-            new_state[T, 2] = self.state[B, 2]
-            new_state[T, 3] = self.state[B, 3]
-            new_state[B, 2] = self.state[F, 2]
-            new_state[B, 3] = self.state[F, 3]
+            new_state[U, 2] = self.state[B, 2]
+            new_state[U, 3] = self.state[B, 3]
+            new_state[D, 2] = self.state[F, 2]
+            new_state[D, 3] = self.state[F, 3]
         # Left Clockwise
-        if action == 6:
-            new_state[F, 0] = self.state[T, 1]
-            new_state[F, 1] = self.state[T, 0]
-            new_state[B, 0] = self.state[B, 1]
-            new_state[B, 1] = self.state[B, 0]
+        elif action == 6:
+            new_state[F, 0] = self.state[U, 1]
+            new_state[F, 1] = self.state[U, 0]
+            new_state[B, 0] = self.state[D, 1]
+            new_state[B, 1] = self.state[D, 0]
             new_state[L, 0] = self.state[L, 1]
             new_state[L, 1] = self.state[L, 3]
             new_state[L, 2] = self.state[L, 0]
             new_state[L, 3] = self.state[L, 2]
-            new_state[T, 0] = self.state[B, 0]
-            new_state[T, 1] = self.state[B, 1]
-            new_state[B, 0] = self.state[F, 0]
-            new_state[B, 1] = self.state[F, 1]
+            new_state[U, 0] = self.state[B, 0]
+            new_state[U, 1] = self.state[B, 1]
+            new_state[D, 0] = self.state[F, 0]
+            new_state[D, 1] = self.state[F, 1]
         # Left Counter-Clockwise
-        if action == 7:
-            new_state[F, 0] = self.state[B, 0]
-            new_state[F, 1] = self.state[B, 1]
-            new_state[B, 0] = self.state[T, 0]
-            new_state[B, 1] = self.state[T, 1]
+        elif action == 7:
+            new_state[F, 0] = self.state[D, 0]
+            new_state[F, 1] = self.state[D, 1]
+            new_state[B, 0] = self.state[U, 0]
+            new_state[B, 1] = self.state[U, 1]
             new_state[L, 0] = self.state[L, 2]
             new_state[L, 1] = self.state[L, 0]
             new_state[L, 2] = self.state[L, 3]
             new_state[L, 3] = self.state[L, 1]
-            new_state[T, 0] = self.state[F, 1]
-            new_state[T, 1] = self.state[F, 0]
-            new_state[B, 0] = self.state[B, 1]
-            new_state[B, 1] = self.state[B, 0]
-        # Top Clockwise
-        if action == 8:
+            new_state[U, 0] = self.state[F, 1]
+            new_state[U, 1] = self.state[F, 0]
+            new_state[D, 0] = self.state[B, 1]
+            new_state[D, 1] = self.state[B, 0]
+        # Up Clockwise
+        elif action == 8:
             new_state[F, 1] = self.state[R, 3]
             new_state[F, 3] = self.state[R, 2]
             new_state[B, 1] = self.state[L, 3]
@@ -158,12 +158,12 @@ class Cube2(gym.Env):
             new_state[R, 3] = self.state[B, 3]
             new_state[L, 2] = self.state[F, 1]
             new_state[L, 3] = self.state[F, 3]
-            new_state[T, 0] = self.state[T, 1]
-            new_state[T, 1] = self.state[T, 3]
-            new_state[T, 2] = self.state[T, 0]
-            new_state[T, 3] = self.state[T, 2]
+            new_state[U, 0] = self.state[U, 1]
+            new_state[U, 1] = self.state[U, 3]
+            new_state[U, 2] = self.state[U, 0]
+            new_state[U, 3] = self.state[U, 2]
         # Up Counter-Clockwise
-        if action == 9:
+        elif action == 9:
             new_state[F, 1] = self.state[L, 2]
             new_state[F, 3] = self.state[L, 3]
             new_state[B, 1] = self.state[R, 2]
@@ -172,12 +172,12 @@ class Cube2(gym.Env):
             new_state[R, 3] = self.state[F, 1]
             new_state[L, 2] = self.state[B, 3]
             new_state[L, 3] = self.state[B, 1]
-            new_state[T, 0] = self.state[T, 2]
-            new_state[T, 1] = self.state[T, 0]
-            new_state[T, 2] = self.state[T, 3]
-            new_state[T, 3] = self.state[T, 1]
+            new_state[U, 0] = self.state[U, 2]
+            new_state[U, 1] = self.state[U, 0]
+            new_state[U, 2] = self.state[U, 3]
+            new_state[U, 3] = self.state[U, 1]
         # Bottom Clockwise
-        if action == 10:
+        elif action == 10:
             new_state[F, 0] = self.state[L, 0]
             new_state[F, 2] = self.state[L, 1]
             new_state[B, 0] = self.state[R, 0]
@@ -186,12 +186,12 @@ class Cube2(gym.Env):
             new_state[R, 1] = self.state[F, 0]
             new_state[L, 0] = self.state[B, 2]
             new_state[L, 1] = self.state[B, 0]
-            new_state[B, 0] = self.state[B, 2]
-            new_state[B, 1] = self.state[B, 0]
-            new_state[B, 2] = self.state[B, 3]
-            new_state[B, 3] = self.state[B, 1]
-        # Down Counter-Clockwise
-        if action == 11:
+            new_state[D, 0] = self.state[D, 2]
+            new_state[D, 1] = self.state[D, 0]
+            new_state[D, 2] = self.state[D, 3]
+            new_state[D, 3] = self.state[D, 1]
+        # Bottom Counter-Clockwise
+        elif action == 11:
             new_state[F, 0] = self.state[R, 1]
             new_state[F, 2] = self.state[R, 0]
             new_state[B, 0] = self.state[L, 1]
@@ -200,23 +200,21 @@ class Cube2(gym.Env):
             new_state[R, 1] = self.state[B, 2]
             new_state[L, 0] = self.state[F, 0]
             new_state[L, 1] = self.state[F, 2]
-            new_state[B, 0] = self.state[B, 1]
-            new_state[B, 1] = self.state[B, 3]
-            new_state[B, 2] = self.state[B, 0]
-            new_state[B, 3] = self.state[B, 2]
-
+            new_state[D, 0] = self.state[D, 1]
+            new_state[D, 1] = self.state[D, 3]
+            new_state[D, 2] = self.state[D, 0]
+            new_state[D, 3] = self.state[D, 2]
         self.state = new_state
-        return self.state, 1 if self._is_solved() else -1, self._is_solved(), self.step_count >= 100, {}
+        return self._get_obs(), 1 if self._is_solved() else -1, self._is_solved(), self.step_count >= 100, {}
 
     def _get_obs(self):
         one_hots = []
         for i in range(6):
-            for j in range(2):
-                for k in range(2):
-                    label = int(self.state[i, j, k])
-                    zeros = np.zeros(6)
-                    zeros[label] = 1
-                    one_hots.append(zeros)
+            for j in range(4):
+                label = int(self.state[i, j])
+                zeros = np.zeros(6)
+                zeros[label] = 1
+                one_hots.append(zeros)
         return np.array(one_hots)
     
     def _is_solved(self):