use notebook

rlcube/cube2.ipynb

@@ -0,0 +1,319 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "dff864f2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gymnasium as gym\n",
+    "import numpy as np\n",
+    "\n",
+    "F = 0\n",
+    "B = 1\n",
+    "R = 2\n",
+    "L = 3\n",
+    "U = 4\n",
+    "D = 5\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.state = np.zeros((6, 4))\n",
+    "        self.step_count = 0\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[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",
+    "    def step(self, action):\n",
+    "        self.step_count += 1\n",
+    "        new_state = self.state.copy()\n",
+    "\n",
+    "        # Front Clockwise\n",
+    "        if action == 0:\n",
+    "            new_state[F, 0] = self.state[F, 2]\n",
+    "            new_state[F, 1] = self.state[F, 0]\n",
+    "            new_state[F, 2] = self.state[F, 3]\n",
+    "            new_state[F, 3] = self.state[F, 1]\n",
+    "            new_state[R, 1] = self.state[U, 3]\n",
+    "            new_state[R, 3] = self.state[U, 1]\n",
+    "            new_state[L, 1] = self.state[D, 3]\n",
+    "            new_state[L, 3] = self.state[D, 1]\n",
+    "            new_state[U, 1] = self.state[L, 1]\n",
+    "            new_state[U, 3] = self.state[L, 3]\n",
+    "            new_state[D, 1] = self.state[R, 1]\n",
+    "            new_state[D, 3] = self.state[R, 3]\n",
+    "        # Front Counter-Clockwise\n",
+    "        elif action == 1:\n",
+    "            new_state[F, 0] = self.state[F, 1]\n",
+    "            new_state[F, 1] = self.state[F, 3]\n",
+    "            new_state[F, 2] = self.state[F, 0]\n",
+    "            new_state[F, 3] = self.state[F, 2]\n",
+    "            new_state[R, 1] = self.state[D, 1]\n",
+    "            new_state[R, 3] = self.state[D, 3]\n",
+    "            new_state[L, 1] = self.state[U, 1]\n",
+    "            new_state[L, 3] = self.state[U, 3]\n",
+    "            new_state[U, 1] = self.state[R, 3]\n",
+    "            new_state[U, 3] = self.state[R, 1]\n",
+    "            new_state[D, 1] = self.state[L, 3]\n",
+    "            new_state[D, 3] = self.state[L, 1]\n",
+    "        # Back Clockwise\n",
+    "        elif action == 2:\n",
+    "            new_state[B, 0] = self.state[B, 1]\n",
+    "            new_state[B, 1] = self.state[B, 3]\n",
+    "            new_state[B, 2] = self.state[B, 0]\n",
+    "            new_state[B, 3] = self.state[B, 2]\n",
+    "            new_state[R, 0] = self.state[D, 0]\n",
+    "            new_state[R, 2] = self.state[D, 2]\n",
+    "            new_state[L, 0] = self.state[U, 0]\n",
+    "            new_state[L, 2] = self.state[U, 2]\n",
+    "            new_state[U, 0] = self.state[R, 2]\n",
+    "            new_state[U, 2] = self.state[R, 0]\n",
+    "            new_state[D, 0] = self.state[L, 2]\n",
+    "            new_state[D, 2] = self.state[L, 0]\n",
+    "        # Back Counter-Clockwise\n",
+    "        elif action == 3:\n",
+    "            new_state[B, 0] = self.state[B, 2]\n",
+    "            new_state[B, 1] = self.state[B, 0]\n",
+    "            new_state[B, 2] = self.state[B, 3]\n",
+    "            new_state[B, 3] = self.state[B, 1]\n",
+    "            new_state[R, 0] = self.state[U, 2]\n",
+    "            new_state[R, 2] = self.state[U, 0]\n",
+    "            new_state[L, 0] = self.state[D, 2]\n",
+    "            new_state[L, 2] = self.state[D, 0]\n",
+    "            new_state[U, 0] = self.state[L, 0]\n",
+    "            new_state[U, 2] = self.state[L, 2]\n",
+    "            new_state[D, 0] = self.state[R, 0]\n",
+    "            new_state[D, 2] = self.state[R, 2]\n",
+    "        # Right Clockwise\n",
+    "        elif action == 4:\n",
+    "            new_state[F, 2] = self.state[D, 2]\n",
+    "            new_state[F, 3] = self.state[D, 3]\n",
+    "            new_state[B, 2] = self.state[U, 2]\n",
+    "            new_state[B, 3] = self.state[U, 3]\n",
+    "            new_state[R, 0] = self.state[R, 2]\n",
+    "            new_state[R, 1] = self.state[R, 0]\n",
+    "            new_state[R, 2] = self.state[R, 3]\n",
+    "            new_state[R, 3] = self.state[R, 1]\n",
+    "            new_state[U, 2] = self.state[F, 3]\n",
+    "            new_state[U, 3] = self.state[F, 2]\n",
+    "            new_state[D, 2] = self.state[B, 3]\n",
+    "            new_state[D, 3] = self.state[B, 2]\n",
+    "        # Right Counter-Clockwise\n",
+    "        elif action == 5:\n",
+    "            new_state[F, 2] = self.state[U, 3]\n",
+    "            new_state[F, 3] = self.state[U, 2]\n",
+    "            new_state[B, 2] = self.state[D, 3]\n",
+    "            new_state[B, 3] = self.state[D, 2]\n",
+    "            new_state[R, 0] = self.state[R, 1]\n",
+    "            new_state[R, 1] = self.state[R, 3]\n",
+    "            new_state[R, 2] = self.state[R, 0]\n",
+    "            new_state[R, 3] = self.state[R, 2]\n",
+    "            new_state[U, 2] = self.state[B, 2]\n",
+    "            new_state[U, 3] = self.state[B, 3]\n",
+    "            new_state[D, 2] = self.state[F, 2]\n",
+    "            new_state[D, 3] = self.state[F, 3]\n",
+    "        # Left Clockwise\n",
+    "        elif action == 6:\n",
+    "            new_state[F, 0] = self.state[U, 1]\n",
+    "            new_state[F, 1] = self.state[U, 0]\n",
+    "            new_state[B, 0] = self.state[D, 1]\n",
+    "            new_state[B, 1] = self.state[D, 0]\n",
+    "            new_state[L, 0] = self.state[L, 1]\n",
+    "            new_state[L, 1] = self.state[L, 3]\n",
+    "            new_state[L, 2] = self.state[L, 0]\n",
+    "            new_state[L, 3] = self.state[L, 2]\n",
+    "            new_state[U, 0] = self.state[B, 0]\n",
+    "            new_state[U, 1] = self.state[B, 1]\n",
+    "            new_state[D, 0] = self.state[F, 0]\n",
+    "            new_state[D, 1] = self.state[F, 1]\n",
+    "        # Left Counter-Clockwise\n",
+    "        elif action == 7:\n",
+    "            new_state[F, 0] = self.state[D, 0]\n",
+    "            new_state[F, 1] = self.state[D, 1]\n",
+    "            new_state[B, 0] = self.state[U, 0]\n",
+    "            new_state[B, 1] = self.state[U, 1]\n",
+    "            new_state[L, 0] = self.state[L, 2]\n",
+    "            new_state[L, 1] = self.state[L, 0]\n",
+    "            new_state[L, 2] = self.state[L, 3]\n",
+    "            new_state[L, 3] = self.state[L, 1]\n",
+    "            new_state[U, 0] = self.state[F, 1]\n",
+    "            new_state[U, 1] = self.state[F, 0]\n",
+    "            new_state[D, 0] = self.state[B, 1]\n",
+    "            new_state[D, 1] = self.state[B, 0]\n",
+    "        # Up Clockwise\n",
+    "        elif action == 8:\n",
+    "            new_state[F, 1] = self.state[R, 3]\n",
+    "            new_state[F, 3] = self.state[R, 2]\n",
+    "            new_state[B, 1] = self.state[L, 3]\n",
+    "            new_state[B, 3] = self.state[L, 2]\n",
+    "            new_state[R, 2] = self.state[B, 1]\n",
+    "            new_state[R, 3] = self.state[B, 3]\n",
+    "            new_state[L, 2] = self.state[F, 1]\n",
+    "            new_state[L, 3] = self.state[F, 3]\n",
+    "            new_state[U, 0] = self.state[U, 1]\n",
+    "            new_state[U, 1] = self.state[U, 3]\n",
+    "            new_state[U, 2] = self.state[U, 0]\n",
+    "            new_state[U, 3] = self.state[U, 2]\n",
+    "        # Up Counter-Clockwise\n",
+    "        elif action == 9:\n",
+    "            new_state[F, 1] = self.state[L, 2]\n",
+    "            new_state[F, 3] = self.state[L, 3]\n",
+    "            new_state[B, 1] = self.state[R, 2]\n",
+    "            new_state[B, 3] = self.state[R, 3]\n",
+    "            new_state[R, 2] = self.state[F, 3]\n",
+    "            new_state[R, 3] = self.state[F, 1]\n",
+    "            new_state[L, 2] = self.state[B, 3]\n",
+    "            new_state[L, 3] = self.state[B, 1]\n",
+    "            new_state[U, 0] = self.state[U, 2]\n",
+    "            new_state[U, 1] = self.state[U, 0]\n",
+    "            new_state[U, 2] = self.state[U, 3]\n",
+    "            new_state[U, 3] = self.state[U, 1]\n",
+    "        # Bottom Clockwise\n",
+    "        elif action == 10:\n",
+    "            new_state[F, 0] = self.state[L, 0]\n",
+    "            new_state[F, 2] = self.state[L, 1]\n",
+    "            new_state[B, 0] = self.state[R, 0]\n",
+    "            new_state[B, 2] = self.state[R, 1]\n",
+    "            new_state[R, 0] = self.state[F, 2]\n",
+    "            new_state[R, 1] = self.state[F, 0]\n",
+    "            new_state[L, 0] = self.state[B, 2]\n",
+    "            new_state[L, 1] = self.state[B, 0]\n",
+    "            new_state[D, 0] = self.state[D, 2]\n",
+    "            new_state[D, 1] = self.state[D, 0]\n",
+    "            new_state[D, 2] = self.state[D, 3]\n",
+    "            new_state[D, 3] = self.state[D, 1]\n",
+    "        # Bottom Counter-Clockwise\n",
+    "        elif action == 11:\n",
+    "            new_state[F, 0] = self.state[R, 1]\n",
+    "            new_state[F, 2] = self.state[R, 0]\n",
+    "            new_state[B, 0] = self.state[L, 1]\n",
+    "            new_state[B, 2] = self.state[L, 0]\n",
+    "            new_state[R, 0] = self.state[B, 0]\n",
+    "            new_state[R, 1] = self.state[B, 2]\n",
+    "            new_state[L, 0] = self.state[F, 0]\n",
+    "            new_state[L, 1] = self.state[F, 2]\n",
+    "            new_state[D, 0] = self.state[D, 1]\n",
+    "            new_state[D, 1] = self.state[D, 3]\n",
+    "            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",
+    "\n",
+    "    def _get_obs(self):\n",
+    "        one_hots = []\n",
+    "        for i in range(6):\n",
+    "            for j in range(4):\n",
+    "                label = int(self.state[i, j])\n",
+    "                zeros = np.zeros(6)\n",
+    "                zeros[label] = 1\n",
+    "                one_hots.append(zeros)\n",
+    "        return np.array(one_hots)\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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "624c83c1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class RewardWrapper(gym.Wrapper):\n",
+    "    def __init__(self, *args, **kwargs):\n",
+    "        super().__init__(*args, **kwargs)\n",
+    "\n",
+    "    def state(self):\n",
+    "        return self.env.state\n",
+    "    \n",
+    "    def reset(self, *args, **kwargs):\n",
+    "        super().reset(*args, **kwargs)\n",
+    "        actions = [self.env.action_space.sample() for _ in range(20)]\n",
+    "        for action in actions:\n",
+    "            self.env.step(action)\n",
+    "        return self.env._get_obs(), {}\n",
+    "\n",
+    "    def step(self, action):\n",
+    "        obs, reward, terminated, truncated, _ = super().step(action)\n",
+    "        return obs, reward, terminated, truncated, _"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "639f54c6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[[1. 1. 0. 3.]\n",
+      " [5. 4. 4. 2.]\n",
+      " [3. 4. 5. 5.]\n",
+      " [1. 2. 2. 4.]\n",
+      " [1. 3. 0. 0.]\n",
+      " [3. 5. 0. 2.]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "env = Cube2()\n",
+    "env = RewardWrapper(env)\n",
+    "obs, _ = env.reset()\n",
+    "print(env.state())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f8b4d968",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from stable_baselines3 import DQN\n",
+    "\n",
+    "model = DQN(\"MlpPolicy\", env, verbose=1)\n",
+    "model.learn(total_timesteps=10000, log_interval=10)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "dev",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

rlcube/main.py

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

rlcube/rlcube/cube2.py

@@ -1,21 +0,0 @@
-import gymnasium as gym
-from .envs.cube2 import Cube2
-from stable_baselines3 import DQN
-
-class RewardWrapper(gym.Wrapper):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-    def step(self, action):
-        obs, reward, terminated, truncated, _ = super().step(action)
-        return obs, reward, terminated, truncated, _
-
-
-def train():
-    env = Cube2()
-    env = RewardWrapper(env)
-
-    model = DQN("MlpPolicy", env, verbose=1)
-    model.learn(total_timesteps=10000, log_interval=10)
-
-    env.close()

rlcube/rlcube/envs/cube2.py

@@ -1,224 +0,0 @@
-from random import shuffle
-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
-        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(), {}
-    
-    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