use adam

rlcube/cube2.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 60,
+   "execution_count": 1,
    "id": "624c83c1",
    "metadata": {},
    "outputs": [
@@ -32,14 +32,14 @@
        ")"
       ]
      },
-     "execution_count": 60,
+     "execution_count": 1,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "from rlcube.models.models import DNN\n",
-    "from rlcube.envs.cube2 import Cube2\n",
+    "from rlcube.envs.cube2 import Cube2Env\n",
     "import numpy as np\n",
     "import torch\n",
     "\n",
@@ -50,7 +50,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 61,
+   "execution_count": 9,
    "id": "16736f3a",
    "metadata": {},
    "outputs": [
@@ -58,32 +58,73 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "tensor([[ 0.0166],\n",
-      "        [ 1.0147],\n",
-      "        [ 1.1610],\n",
-      "        [ 0.9844],\n",
-      "        [-0.0268],\n",
-      "        [ 1.1526]], grad_fn=<AddmmBackward0>)\n",
-      "tensor([10,  1,  5,  0, 10,  1])\n"
+      "rotationController.setState([[0, 0, 4, 4], [1, 1, 5, 5], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 1, 1], [5, 5, 0, 0]]);\n",
+      "0.40487873554229736\n",
+      "4\n",
+      "\n",
+      "rotationController.setState([[0, 4, 0, 4], [1, 1, 5, 5], [2, 5, 2, 0], [3, 4, 3, 1], [4, 2, 1, 2], [5, 3, 0, 3]]);\n",
+      "0.0839405208826065\n",
+      "7\n",
+      "\n",
+      "rotationController.setState([[0, 4, 0, 4], [5, 1, 5, 1], [1, 5, 4, 0], [0, 4, 5, 1], [3, 2, 3, 2], [2, 3, 2, 3]]);\n",
+      "-0.23320673406124115\n",
+      "3\n",
+      "\n",
+      "rotationController.setState([[0, 5, 0, 1], [5, 4, 5, 0], [1, 5, 4, 4], [0, 4, 1, 1], [3, 3, 2, 2], [2, 3, 2, 3]]);\n",
+      "0.31869572401046753\n",
+      "0\n",
+      "\n",
+      "rotationController.setState([[5, 5, 1, 1], [4, 4, 0, 0], [5, 5, 4, 4], [0, 0, 1, 1], [3, 3, 2, 2], [3, 3, 2, 2]]);\n",
+      "-0.16905824840068817\n",
+      "7\n",
+      "\n",
+      "rotationController.setState([[5, 4, 1, 4], [4, 1, 0, 1], [5, 5, 4, 0], [0, 0, 5, 1], [3, 2, 3, 2], [3, 3, 2, 2]]);\n",
+      "0.20266102254390717\n",
+      "3\n",
+      "\n",
+      "rotationController.setState([[2, 3, 1, 4], [3, 3, 0, 1], [5, 5, 4, 0], [0, 1, 0, 5], [4, 1, 3, 2], [5, 4, 2, 2]]);\n",
+      "0.6111429333686829\n",
+      "3\n",
+      "\n",
+      "rotationController.setState([[2, 0, 1, 4], [3, 5, 0, 0], [5, 5, 3, 1], [0, 1, 3, 4], [1, 2, 4, 3], [5, 4, 2, 2]]);\n",
+      "1.3550236225128174\n",
+      "2\n",
+      "\n",
+      "rotationController.setState([[0, 0, 1, 4], [5, 5, 5, 0], [1, 2, 3, 1], [0, 3, 3, 4], [1, 2, 4, 3], [2, 5, 2, 4]]);\n",
+      "0.9975889325141907\n",
+      "7\n",
+      "\n",
+      "rotationController.setState([[2, 0, 1, 4], [3, 5, 0, 0], [5, 5, 3, 1], [0, 1, 3, 4], [1, 2, 4, 3], [5, 4, 2, 2]]);\n",
+      "1.3550236225128174\n",
+      "2\n",
+      "\n"
      ]
     }
    ],
    "source": [
-    "env = Cube2()\n",
-    "obs, _ = env.reset()\n",
-    "obs1, _, _, _, _ = env.step(0)\n",
-    "obs2, _, _, _, _ = env.step(0)\n",
-    "obs3, _, _, _, _ = env.step(2)\n",
-    "obs4, _, _, _, _ = env.step(2)\n",
+    "batch_obs = []\n",
+    "env = Cube2Env()\n",
     "for _ in range(10):\n",
-    "    obsMany, _, _, _, _ = env.step(env.action_space.sample())\n",
-    "batched_obs = torch.tensor(\n",
-    "    np.array([obs, obs1, obs2, obs3, obs4, obsMany]), dtype=torch.float32\n",
-    ")\n",
+    "    obs, _, _, _, _ = env.step(env.action_space.sample())\n",
+    "    batch_obs.append(torch.tensor(obs, dtype=torch.float32))\n",
+    "batched_obs = torch.stack(batch_obs)\n",
     "out = net(batched_obs)\n",
-    "print(out[\"value\"])\n",
-    "print(torch.argmax(out[\"policy\"], dim=1))"
+    "\n",
+    "for i in range(10):\n",
+    "    env = Cube2Env.from_obs(batch_obs[i])\n",
+    "    env.print_js_code()\n",
+    "    print(out[\"value\"][i].item())\n",
+    "    print(torch.argmax(out[\"policy\"][i]).item())\n",
+    "    print()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aee2a911",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

rlcube/rlcube/train/train.py

@@ -17,7 +17,7 @@ print(f"Using device: {device}")
 
 def train(epochs: int = 100):
     if not os.path.exists("dataset.pt"):
-        create_dataset(num_envs=10000, num_steps=20, filepath="dataset.pt")
+        create_dataset(num_envs=1000, num_steps=20, filepath="dataset.pt")
     dataset = Cube2Dataset("dataset.pt")
     print("Number of samples:", len(dataset))
     print("Number of epochs:", epochs)
@@ -29,7 +29,7 @@ def train(epochs: int = 100):
     if os.path.exists("models/model_best.pth"):
         net.load("models/model_best.pth")
     net = net.to(device)
-    optimizer = torch.optim.RMSprop(net.parameters(), lr=0.000001)
+    optimizer = torch.optim.Adam(net.parameters(), lr=0.000001)
     value_loss_fn = torch.nn.MSELoss(reduction="none")
     policy_loss_fn = torch.nn.CrossEntropyLoss(reduction="none")