train

rlcube/cube2.ipynb

@@ -2,14 +2,85 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 60,
    "id": "624c83c1",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "DNN(\n",
+       "  (fc_in): Linear(in_features=144, out_features=512, bias=True)\n",
+       "  (residual_blocks): ModuleList(\n",
+       "    (0-3): 4 x ResidualBlock(\n",
+       "      (ln1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
+       "      (fc1): Linear(in_features=512, out_features=1024, bias=True)\n",
+       "      (ln2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
+       "      (fc2): Linear(in_features=1024, out_features=512, bias=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (fc_value): Sequential(\n",
+       "    (0): Linear(in_features=512, out_features=64, bias=True)\n",
+       "    (1): ReLU()\n",
+       "    (2): Linear(in_features=64, out_features=1, bias=True)\n",
+       "  )\n",
+       "  (fc_policy): Sequential(\n",
+       "    (0): Linear(in_features=512, out_features=64, bias=True)\n",
+       "    (1): ReLU()\n",
+       "    (2): Linear(in_features=64, out_features=12, bias=True)\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 60,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "from rlcube.train.train import train\n",
+    "from rlcube.models.models import DNN\n",
+    "from rlcube.envs.cube2 import Cube2\n",
+    "import numpy as np\n",
+    "import torch\n",
     "\n",
-    "train()"
+    "net = DNN()\n",
+    "net.load(\"models/model_best.pth\")\n",
+    "net.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 61,
+   "id": "16736f3a",
+   "metadata": {},
+   "outputs": [
+    {
+     "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"
+     ]
+    }
+   ],
+   "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",
+    "for _ in range(10):\n",
+    "    obsMany, _, _, _, _ = env.step(env.action_space.sample())\n",
+    "batched_obs = torch.tensor(np.array([obs, obs1, obs2, obs3, obs4, obsMany]), dtype=torch.float32)\n",
+    "out = net(batched_obs)\n",
+    "print(out[\"value\"])\n",
+    "print(torch.argmax(out[\"policy\"], dim=1))"
    ]
   }
  ],

rlcube/rlcube/train/train.py

@@ -29,9 +29,9 @@ 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.0001)
-    value_loss_fn = torch.nn.MSELoss()
-    policy_loss_fn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.RMSprop(net.parameters(), lr=0.000001)
+    value_loss_fn = torch.nn.MSELoss(reduction="none")
+    policy_loss_fn = torch.nn.CrossEntropyLoss(reduction="none")
 
     best_loss = float("inf")
     for epoch in range(epochs):
@@ -56,9 +56,9 @@ def train(epochs: int = 100):
             target_values, indices = (neighbors_values + neighbors_rewards).max(dim=1)
             indices = indices.reshape(-1)
 
-            loss_v = value_loss_fn(values, target_values)
-            loss_p = policy_loss_fn(policies, indices)
-            loss = loss_v + loss_p
+            loss_v = value_loss_fn(values, target_values).reshape(-1) / D.reshape(-1).detach()
+            loss_p = policy_loss_fn(policies, indices).reshape(-1) / D.reshape(-1).detach()
+            loss = (loss_v + loss_p).mean()
             epoch_loss += loss.item()
             optimizer.zero_grad()
             loss.backward()