update notebook

alex-petrenko · Jan 26, 2025 · 4b7d189 · 4b7d189
1 parent d84bafe
commit 4b7d189
Showing 1 changed file with 30 additions and 18 deletions.
diff --git a/sf_examples/nethack/notebooks/simba_sizes.ipynb b/sf_examples/nethack/notebooks/simba_sizes.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -17,24 +17,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[33m[2025-01-26 08:25:41,847][2216020] Environment nethack_progress already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,848][2216020] Environment nethack_staircase already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,848][2216020] Environment nethack_score already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,849][2216020] Environment nethack_pet already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,850][2216020] Environment nethack_oracle already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,851][2216020] Environment nethack_gold already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,851][2216020] Environment nethack_eat already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,852][2216020] Environment nethack_scout already registered, overwriting...\u001b[0m\n",
-      "\u001b[33m[2025-01-26 08:25:41,852][2216020] Environment nethack_challenge already registered, overwriting...\u001b[0m\n",
-      "\u001b[36m[2025-01-26 08:25:41,853][2216020] register_encoder_factory: <function make_nethack_encoder at 0x7a184aca7b50>\u001b[0m\n",
-      "\u001b[36m[2025-01-26 08:25:41,853][2216020] register_actor_critic_factory: <function make_nethack_actor_critic at 0x7a184aca7a30>\u001b[0m\n"
+      "\u001b[36m[2025-01-26 09:32:39,142][825930] register_encoder_factory: <function make_nethack_encoder at 0x71519d7113f0>\u001b[0m\n",
+      "\u001b[36m[2025-01-26 09:32:39,144][825930] register_actor_critic_factory: <function make_nethack_actor_critic at 0x71519d7112d0>\u001b[0m\n"
      ]
     }
    ],
@@ -48,7 +39,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -57,10 +48,10 @@
      "text": [
       "Hidden Dim    16     32     64     128     256     512\n",
       "Depth                                                 \n",
-      "1           0.38M  0.77M  1.58M  3.35M   7.47M  18.08M\n",
-      "2           0.39M  0.79M  1.65M  3.64M   8.65M  22.80M\n",
-      "3           0.39M  0.81M  1.73M  3.94M   9.83M  27.52M\n",
-      "4           0.40M  0.83M  1.80M  4.23M  11.02M  32.24M\n"
+      "1           0.39M  0.78M  1.59M  3.35M   7.48M  18.09M\n",
+      "2           0.39M  0.79M  1.66M  3.65M   8.66M  22.81M\n",
+      "3           0.40M  0.81M  1.73M  3.94M   9.84M  27.53M\n",
+      "4           0.40M  0.83M  1.81M  4.24M  11.02M  32.25M\n"
      ]
     }
    ],
@@ -87,6 +78,27 @@
     "print(pivot_table)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model Size: 3.28M\n"
+     ]
+    }
+   ],
+   "source": [
+    "from sf_examples.nethack.models import ChaoticDwarvenGPT5\n",
+    "\n",
+    "model = ChaoticDwarvenGPT5(cfg, env_info.obs_space)\n",
+    "total_params = sum(p.numel() for p in model.parameters())\n",
+    "print(f\"Model Size: {total_params / 10**6:.2f}M\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,