add cube2 env

package.json

@@ -6,8 +6,8 @@
     "dev": "next dev --turbopack",
     "build": "next build --turbopack",
     "start": "next start",
-    "lint": "eslint",
-    "format": "prettier --write ."
+    "lint": "eslint && cd rlcube && uvx ruff check",
+    "format": "prettier --write . && cd rlcube && uvx ruff format"
   },
   "dependencies": {
     "@heroui/react": "^2.8.3",

rlcube/cube2.ipynb

@@ -17,26 +17,29 @@
     "U = 4\n",
     "D = 5\n",
     "\n",
+    "\n",
     "class Cube2(gym.Env):\n",
     "    def __init__(self):\n",
     "        super().__init__()\n",
     "        self.action_space = gym.spaces.Discrete(12)\n",
-    "        self.observation_space = gym.spaces.Box(low=0,high=1,shape=(24, 6),dtype=np.int8)\n",
+    "        self.observation_space = gym.spaces.Box(\n",
+    "            low=0, high=1, shape=(24, 6), dtype=np.int8\n",
+    "        )\n",
     "        self.state = np.zeros((6, 4))\n",
     "        self.step_count = 0\n",
-    "    \n",
+    "\n",
     "    def reset(self, seed=None, options=None):\n",
     "        super().reset(seed=seed, options=options)\n",
     "        self.state = np.zeros((6, 4))\n",
     "        self.state[0] = np.ones(4) * F\n",
     "        self.state[1] = np.ones(4) * B\n",
     "        self.state[2] = np.ones(4) * R\n",
-    "        self.state[3] = np.ones(4) * L \n",
+    "        self.state[3] = np.ones(4) * L\n",
     "        self.state[4] = np.ones(4) * U\n",
     "        self.state[5] = np.ones(4) * D\n",
     "        self.step_count = 0\n",
     "        return self._get_obs(), {}\n",
-    "    \n",
+    "\n",
     "    def step(self, action):\n",
     "        self.step_count += 1\n",
     "        new_state = self.state.copy()\n",
@@ -210,7 +213,13 @@
     "            new_state[D, 2] = self.state[D, 0]\n",
     "            new_state[D, 3] = self.state[D, 2]\n",
     "        self.state = new_state\n",
-    "        return self._get_obs(), 1 if self._is_solved() else -1, self._is_solved(), self.step_count >= 100, {}\n",
+    "        return (\n",
+    "            self._get_obs(),\n",
+    "            1 if self._is_solved() else -1,\n",
+    "            self._is_solved(),\n",
+    "            self.step_count >= 100,\n",
+    "            {},\n",
+    "        )\n",
     "\n",
     "    def _get_obs(self):\n",
     "        one_hots = []\n",
@@ -221,12 +230,12 @@
     "                zeros[label] = 1\n",
     "                one_hots.append(zeros)\n",
     "        return np.array(one_hots)\n",
-    "    \n",
+    "\n",
     "    def _is_solved(self):\n",
     "        for i in range(6):\n",
     "            if np.mean(self.state[i]) != self.state[i][0]:\n",
     "                return False\n",
-    "        return True\n"
+    "        return True"
    ]
   },
   {
@@ -242,10 +251,10 @@
     "\n",
     "    def state(self):\n",
     "        return self.env.state\n",
-    "    \n",
+    "\n",
     "    def step_count(self):\n",
     "        return self.env.step_count\n",
-    "    \n",
+    "\n",
     "    def reset(self, *args, **kwargs):\n",
     "        self.env.reset(*args, **kwargs)\n",
     "        for _ in range(4):\n",
@@ -501,7 +510,7 @@
     "print(f\"rotationController.addRotationStepCode(...{json.dumps(solved_actions)})\")\n",
     "\n",
     "print()\n",
-    "print(f\"Solved in {len(solved_actions)} steps\")\n"
+    "print(f\"Solved in {len(solved_actions)} steps\")"
    ]
   }
  ],

rlcube/envs/cube2.py

@@ -0,0 +1,229 @@
+import gymnasium as gym
+import numpy as np
+
+F = 0
+B = 1
+R = 2
+L = 3
+U = 4
+D = 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, 4))
+        self.step_count = 0
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed, options=options)
+        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) * U
+        self.state[5] = np.ones(4) * D
+        self.step_count = 0
+        return self._get_obs(), {}
+
+    def step(self, action):
+        self.step_count += 1
+        new_state = self.state.copy()
+
+        # Front Clockwise
+        if action == 0:
+            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[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
+        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[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
+        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[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
+        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[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
+        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[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
+        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[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
+        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[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
+        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[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]
+            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[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
+        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]
+            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[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
+        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]
+            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[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]
+            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[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._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(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):
+        for i in range(6):
+            if np.mean(self.state[i]) != self.state[i][0]:
+                return False
+        return True