auto python state

rlcube/main.py

@@ -1,7 +1,9 @@
-from rlcube.cube2 import train
+from rlcube.envs.cube2 import Cube2
 
 def main():
-    train()
+    env = Cube2()
+    obs, _ = env.reset()
+    print(obs)
 
 
 if __name__ == "__main__":

rlcube/rlcube/envs/cube2.py

@@ -2,35 +2,35 @@ from random import shuffle
 import gymnasium as gym
 import numpy as np
 
-FRONT = 0
-RIGHT = 1
-BACK = 2
-LEFT = 3
-UP = 4
-DOWN = 5
+F = 0
+B = 1
+R = 2
+L = 3
+T = 4
+B = 5
 
 class Cube2(gym.Env):
     def __init__(self):
         super().__init__()
         self.action_space = gym.spaces.Discrete(12)
         self.observation_space = gym.spaces.Box(low=0,high=1,shape=(24, 6),dtype=np.int8)
-        self.state = np.zeros((6, 2, 2))
+        self.state = np.zeros((6, 4))
         self.step_count = 0
     
     def reset(self, seed=None, options=None):
         super().reset(seed=seed, options=options)
-        self.state = np.zeros((6, 2, 2))
-        self.state[0] = np.ones((2, 2)) * FRONT
-        self.state[1] = np.ones((2, 2)) * RIGHT
-        self.state[2] = np.ones((2, 2)) * BACK
-        self.state[3] = np.ones((2, 2)) * LEFT 
-        self.state[4] = np.ones((2, 2)) * UP
-        self.state[5] = np.ones((2, 2)) * DOWN
-        shuffle_steps =self.np_random.integers(0, 20)
-        for i in range(shuffle_steps):
-            self.step(self.action_space.sample())
+        self.state = np.zeros((6, 4))
+        self.state[0] = np.ones(4) * F
+        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.step_count = 0
-        return self._get_obs(), {}
+        return self.state, {}
     
     def step(self, action):
         self.step_count += 1
@@ -38,127 +38,175 @@ class Cube2(gym.Env):
 
         # Front Clockwise
         if action == 0:
-            new_state[RIGHT, 0, 0] = self.state[UP, 0, 0]
-            new_state[RIGHT, 1, 0] = self.state[UP, 1, 0]
-            new_state[DOWN, 0, 1]  = self.state[RIGHT, 1, 0]
-            new_state[DOWN, 1, 1]  = self.state[RIGHT, 0, 0]
-            new_state[LEFT, 0, 1]  = self.state[DOWN, 0, 1]
-            new_state[LEFT, 1, 1]  = self.state[DOWN, 1, 1]
-            new_state[UP, 0, 0]    = self.state[LEFT, 1, 1]
-            new_state[UP, 1, 0]    = self.state[LEFT, 0, 1]
+            new_state[F, 0] = self.state[F, 2]
+            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]
         # Front Counter-Clockwise
         if action == 1:
-            new_state[LEFT, 0, 1]  = self.state[UP, 1, 0]
-            new_state[LEFT, 1, 1]  = self.state[UP, 0, 0]
-            new_state[DOWN, 0, 1]  = self.state[LEFT, 0, 1]
-            new_state[DOWN, 1, 1]  = self.state[LEFT, 1, 1]
-            new_state[RIGHT, 0, 0] = self.state[DOWN, 1, 1]
-            new_state[RIGHT, 1, 0] = self.state[DOWN, 0, 1]
-            new_state[UP, 0, 0]    = self.state[RIGHT, 0, 0]
-            new_state[UP, 1, 0]    = self.state[RIGHT, 1, 0]
-        # Right Clockwise
+            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]
+        # Back Clockwise
         if action == 2:
-            new_state[BACK, 0, 0]  = self.state[UP, 1, 0]
-            new_state[BACK, 1, 0]  = self.state[UP, 1, 1]
-            new_state[DOWN, 1, 0]  = self.state[BACK, 0, 0]
-            new_state[DOWN, 1, 1]  = self.state[BACK, 1, 0]
-            new_state[FRONT, 0, 1] = self.state[DOWN, 1, 1]
-            new_state[FRONT, 1, 1] = self.state[DOWN, 1, 0]
-            new_state[UP, 1, 0]    = self.state[FRONT, 1, 1]
-            new_state[UP, 1, 1]    = self.state[FRONT, 0, 1]
-        # Right Counter-Clockwise
+            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]
+        # Back Counter-Clockwise
         if action == 3:
-            new_state[FRONT, 0, 1] = self.state[UP, 1, 1]
-            new_state[FRONT, 1, 1] = self.state[UP, 1, 0]
-            new_state[DOWN, 1, 1]  = self.state[FRONT, 0, 1]
-            new_state[DOWN, 1, 0]  = self.state[FRONT, 1, 1]
-            new_state[BACK, 0, 0]  = self.state[DOWN, 1, 0]
-            new_state[BACK, 1, 0]  = self.state[DOWN, 1, 1]
-            new_state[UP, 1, 0]    = self.state[BACK, 0, 0]
-            new_state[UP, 1, 1]    = self.state[BACK, 1, 0]
-        # Left Clockwise
+            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]
+        # Right Clockwise
         if action == 4:
-            new_state[FRONT, 0, 0] = self.state[UP, 0, 1]
-            new_state[FRONT, 1, 0] = self.state[UP, 0, 0]
-            new_state[DOWN, 0, 0]  = self.state[FRONT, 1, 0]
-            new_state[DOWN, 0, 1]  = self.state[FRONT, 0, 0]
-            new_state[BACK, 0, 1]  = self.state[DOWN, 0, 0]
-            new_state[BACK, 1, 1]  = self.state[DOWN, 0, 1]
-            new_state[UP, 0, 0]    = self.state[BACK, 0, 1]
-            new_state[UP, 0, 1]    = self.state[BACK, 1, 1]
-        # Left Counter-Clockwise
+            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]
+            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]
+        # Right Counter-Clockwise
         if action == 5:
-            new_state[BACK, 0, 1]  = self.state[UP, 0, 0]
-            new_state[BACK, 1, 1]  = self.state[UP, 0, 1]
-            new_state[DOWN, 0, 0]  = self.state[BACK, 0, 1]
-            new_state[DOWN, 0, 1]  = self.state[BACK, 1, 1]
-            new_state[FRONT, 0, 0] = self.state[DOWN, 0, 1]
-            new_state[FRONT, 1, 0] = self.state[DOWN, 0, 0]
-            new_state[UP, 0, 1]    = self.state[FRONT, 0, 0]
-            new_state[UP, 0, 0]    = self.state[FRONT, 1, 0]
-        # Back Clockwise
+            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]
+            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]
+        # Left Clockwise
         if action == 6:
-            new_state[RIGHT, 0, 1] = self.state[DOWN, 1, 0]
-            new_state[RIGHT, 1, 1] = self.state[DOWN, 0, 0]
-            new_state[UP, 0, 1]    = self.state[RIGHT, 0, 1]
-            new_state[UP, 1, 1]    = self.state[RIGHT, 1, 1]
-            new_state[LEFT, 0, 0]  = self.state[UP, 1, 1]
-            new_state[LEFT, 1, 0]  = self.state[UP, 0, 1]
-            new_state[DOWN, 0, 0]  = self.state[LEFT, 0, 0]
-            new_state[DOWN, 1, 0]  = self.state[LEFT, 1, 0]
-        # Back Counter-Clockwise
+            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]
+            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]
+        # Left Counter-Clockwise
         if action == 7:
-            new_state[DOWN, 1, 0]  = self.state[RIGHT, 0, 1] 
-            new_state[DOWN, 0, 0]  = self.state[RIGHT, 1, 1] 
-            new_state[RIGHT, 0, 1] = self.state[UP, 0, 1]    
-            new_state[RIGHT, 1, 1] = self.state[UP, 1, 1]    
-            new_state[UP, 1, 1]    = self.state[LEFT, 0, 0]  
-            new_state[UP, 0, 1]    = self.state[LEFT, 1, 0]  
-            new_state[LEFT, 0, 0]  = self.state[DOWN, 0, 0]  
-            new_state[LEFT, 1, 0]  = self.state[DOWN, 1, 0]
-        # Up Clockwise
+            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]
+            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[FRONT, 0, 0] = self.state[RIGHT, 0, 0]
-            new_state[FRONT, 0, 1] = self.state[RIGHT, 0, 1]
-            new_state[LEFT, 0, 0]  = self.state[FRONT, 0, 0]
-            new_state[LEFT, 0, 1]  = self.state[FRONT, 0, 1]
-            new_state[BACK, 0, 0]  = self.state[LEFT, 0, 0]
-            new_state[BACK, 0, 1]  = self.state[LEFT, 0, 1]
-            new_state[RIGHT, 0, 0] = self.state[BACK, 0, 0]
-            new_state[RIGHT, 0, 1] = self.state[BACK, 0, 1]
+            new_state[F, 1] = self.state[R, 3]
+            new_state[F, 3] = self.state[R, 2]
+            new_state[B, 1] = self.state[L, 3]
+            new_state[B, 3] = self.state[L, 2]
+            new_state[R, 2] = self.state[B, 1]
+            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]
         # Up Counter-Clockwise
         if action == 9:
-            new_state[RIGHT, 0, 0] = self.state[FRONT, 0, 0]
-            new_state[RIGHT, 0, 1] = self.state[FRONT, 0, 1]
-            new_state[FRONT, 0, 0] = self.state[LEFT, 0, 0] 
-            new_state[FRONT, 0, 1] = self.state[LEFT, 0, 1] 
-            new_state[LEFT, 0, 0]  = self.state[BACK, 0, 0] 
-            new_state[LEFT, 0, 1]  = self.state[BACK, 0, 1] 
-            new_state[BACK, 0, 0]  = self.state[RIGHT, 0, 0]
-            new_state[BACK, 0, 1]  = self.state[RIGHT, 0, 1]
-        # Down Clockwise
+            new_state[F, 1] = self.state[L, 2]
+            new_state[F, 3] = self.state[L, 3]
+            new_state[B, 1] = self.state[R, 2]
+            new_state[B, 3] = self.state[R, 3]
+            new_state[R, 2] = self.state[F, 3]
+            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]
+        # Bottom Clockwise
         if action == 10:
-            new_state[FRONT, 1, 0] = self.state[LEFT, 1, 0]
-            new_state[FRONT, 1, 1] = self.state[LEFT, 1, 1]
-            new_state[LEFT, 1, 0]  = self.state[BACK, 1, 0]
-            new_state[LEFT, 1, 1]  = self.state[BACK, 1, 1]
-            new_state[BACK, 1, 0]  = self.state[RIGHT, 1, 0]
-            new_state[BACK, 1, 1]  = self.state[RIGHT, 1, 1]
-            new_state[RIGHT, 1, 0] = self.state[FRONT, 1, 0]
-            new_state[RIGHT, 1, 1] = self.state[FRONT, 1, 1]
+            new_state[F, 0] = self.state[L, 0]
+            new_state[F, 2] = self.state[L, 1]
+            new_state[B, 0] = self.state[R, 0]
+            new_state[B, 2] = self.state[R, 1]
+            new_state[R, 0] = self.state[F, 2]
+            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[LEFT, 1, 0]  = self.state[FRONT, 1, 0]
-            new_state[LEFT, 1, 1]  = self.state[FRONT, 1, 1]
-            new_state[BACK, 1, 0]  = self.state[LEFT, 1, 0] 
-            new_state[BACK, 1, 1]  = self.state[LEFT, 1, 1] 
-            new_state[RIGHT, 1, 0] = self.state[BACK, 1, 0] 
-            new_state[RIGHT, 1, 1] = self.state[BACK, 1, 1] 
-            new_state[FRONT, 1, 0] = self.state[RIGHT, 1, 0]
-            new_state[FRONT, 1, 1] = self.state[RIGHT, 1, 1]
+            new_state[F, 0] = self.state[R, 1]
+            new_state[F, 2] = self.state[R, 0]
+            new_state[B, 0] = self.state[L, 1]
+            new_state[B, 2] = self.state[L, 0]
+            new_state[R, 0] = self.state[B, 0]
+            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]
 
         self.state = new_state
-        return self._get_obs(), 1 if self._is_solved() else -1, self._is_solved(), self.step_count >= 100, {}
+        return self.state, 1 if self._is_solved() else -1, self._is_solved(), self.step_count >= 100, {}
 
     def _get_obs(self):
         one_hots = []
@@ -173,6 +221,6 @@ class Cube2(gym.Env):
     
     def _is_solved(self):
         for i in range(6):
-            if np.mean(self.state[i]) != self.state[i][0][0]:
+            if np.mean(self.state[i]) != self.state[i][0]:
                 return False
         return True

src/components/state-modal.tsx

@@ -55,12 +55,12 @@ export const StateModal = forwardRef<StateModalRef, unknown>((_, ref) => {
                   <div className="font-mono">[{state[3].map((index) => Index2Color[index]).join(', ')}]</div>
                 </div>
                 <div className="flex gap-2 items-center">
-                  <div className="text-sm w-24 font-mont">Up</div>
+                  <div className="text-sm w-24 font-mont">Top</div>
                   <div className="font-mono">{JSON.stringify(state[4])}</div>
                   <div className="font-mono">[{state[4].map((index) => Index2Color[index]).join(', ')}]</div>
                 </div>
                 <div className="flex gap-2 items-center">
-                  <div className="text-sm w-24 font-mont">Down</div>
+                  <div className="text-sm w-24 font-mont">Bottom</div>
                   <div className="font-mono">{JSON.stringify(state[5])}</div>
                   <div className="font-mono">[{state[5].map((index) => Index2Color[index]).join(', ')}]</div>
                 </div>