Fix volume transformation, add observation wrapper

docs/02_notebooks/L3_slicing.ipynb

@@ -252,13 +252,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from armscan_env.slicing import slice_volume\n",
+    "from armscan_env.envs.state_action import ManipulatorAction\n",
+    "from armscan_env.volumes.slicing import slice_volume\n",
     "\n",
     "sliced_volume = slice_volume(\n",
-    "    z_rotation=19.3,\n",
-    "    x_rotation=0.0,\n",
-    "    x_trans=0.0,\n",
-    "    y_trans=140.0,\n",
+    "    action=ManipulatorAction(rotation=(19.3, 0), translation=(0, 140)),\n",
     "    volume=volume,\n",
     "    slice_shape=(volume.GetSize()[0], volume.GetSize()[2]),\n",
     ")\n",
@@ -323,9 +321,10 @@
     "    sliced_volume = slice_volume(\n",
     "        volume=volume,\n",
     "        slice_shape=(volume.GetSize()[0], volume.GetSize()[2]),\n",
-    "        z_rotation=z[i],\n",
-    "        x_rotation=0,\n",
-    "        y_trans=t[i],\n",
+    "        action=ManipulatorAction(\n",
+    "            rotation=(z[i], 0),\n",
+    "            translation=(0, t[i]),\n",
+    "        ),\n",
     "    )\n",
     "    sliced_img = sitk.GetArrayFromImage(sliced_volume)[:, 0, :]\n",
     "    ax2.set_title(f\"Slice {i}\")\n",

docs/02_notebooks/L4_environment.ipynb

@@ -37,8 +37,8 @@
     "from armscan_env.clustering import TissueClusters\n",
     "from armscan_env.config import get_config\n",
     "from armscan_env.envs.rewards import anatomy_based_rwd\n",
-    "from armscan_env.slicing import slice_volume\n",
     "from armscan_env.util.visualizations import show_clusters\n",
+    "from armscan_env.volumes.slicing import slice_volume\n",
     "from IPython.core.display import HTML\n",
     "\n",
     "config = get_config()"

docs/02_notebooks/L5_linear_sweep.ipynb

@@ -5,18 +5,17 @@
    "execution_count": null,
    "id": "a4e98c0276b6012d",
    "metadata": {},
-   "outputs": [],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "50b440b37fd9414b",
    "metadata": {},
-   "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
@@ -37,7 +36,8 @@
     "from tianshou.highlevel.env import EnvMode\n",
     "\n",
     "config = get_config()"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "markdown",
@@ -52,7 +52,6 @@
    "execution_count": null,
    "id": "9ed46c7b",
    "metadata": {},
-   "outputs": [],
    "source": [
     "def walk_through_env(\n",
     "    env: LabelmapEnv,\n",
@@ -123,27 +122,27 @@
     "\n",
     "    if show:\n",
     "        plt.show()"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "da45ed45bb7b8f3b",
    "metadata": {},
-   "outputs": [],
    "source": [
     "volume_1 = sitk.ReadImage(config.get_labels_path(1))\n",
     "volume_2 = sitk.ReadImage(config.get_labels_path(2))\n",
     "img_array_1 = sitk.GetArrayFromImage(volume_1)\n",
     "img_array_2 = sitk.GetArrayFromImage(volume_2)"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "63dd92db3829d7db",
    "metadata": {},
-   "outputs": [],
    "source": [
     "volume_size = volume_1.GetSize()\n",
     "\n",
@@ -159,36 +158,36 @@
     "    render_mode=\"animation\",\n",
     "    n_stack=2,\n",
     ").create_env(EnvMode.WATCH)"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "16a139f61aaafd19",
    "metadata": {},
-   "outputs": [],
    "source": [
     "env_rollout = walk_through_env(env, 10)\n",
     "\n",
     "plot_rollout_rewards(env_rollout)"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "6cdf855cc85a743a",
    "metadata": {},
-   "outputs": [],
    "source": [
     "env.get_cur_animation_as_html()"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "519dde5f1cea8a5f",
    "metadata": {},
-   "outputs": [],
    "source": [
     "volume_size = volume_1.GetSize()\n",
     "\n",
@@ -206,54 +205,55 @@
     "    project_actions_to=\"y\",\n",
     "    apply_volume_transformation=True,\n",
     ").create_env(EnvMode.WATCH)"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "22877ab71fed2eb0",
    "metadata": {},
-   "outputs": [],
    "source": [
     "projected_env_rollout = walk_through_env(\n",
     "    projected_env,\n",
     "    10,\n",
     "    render_title=\"Projected labelmap slice\",\n",
     ")\n",
     "plot_rollout_rewards(projected_env_rollout)"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "c2779884526e0716",
    "metadata": {},
-   "outputs": [],
    "source": [
     "print(\n",
     "    \"Observed 'rewards': \\n\",\n",
     "    [round(obs[1][-1], 4) for obs in projected_env_rollout.observations],\n",
     ")\n",
     "print(\"Env rewards: \\n\", [round(r, 4) for r in projected_env_rollout.rewards])"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "6ada94c94fe77de0",
    "metadata": {},
-   "outputs": [],
    "source": [
     "projected_env.get_cur_animation_as_html()"
-   ]
+   ],
+   "outputs": []
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "4fb82c1487521b11",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "source": [],
+   "outputs": []
   }
  ],
  "metadata": {

pyproject.toml

@@ -154,7 +154,7 @@ lint = ["_black_check", "_ruff_check", "_ruff_check_nb"]
 _poetry_install_sort_plugin = "poetry self add poetry-plugin-sort"
 _poetry_sort = "poetry sort"
 clean-nbs = "python docs/nbstripout.py"
-format = ["_black_format", "_ruff_format", "_ruff_format_nb", "_poetry_install_sort_plugin", "_poetry_sort"]
+format = ["_black_format", "_ruff_format", "_ruff_format_nb"]
 _autogen_rst = "python docs/autogen_rst.py"
 _sphinx_build = "sphinx-build -W -b html docs docs/_build"
 _jb_generate_toc = "python docs/create_toc.py"

scripts/armscan_array_obs.py

@@ -38,9 +38,9 @@
         num_test_envs=1,
         buffer_size=100000,
         batch_size=256,
-        step_per_collect=10,
-        update_per_step=100,
-        start_timesteps=500,
+        step_per_collect=200,
+        update_per_step=2,
+        start_timesteps=5000,
         start_timesteps_random=True,
     )
 
@@ -61,6 +61,7 @@
         termination_criterion=LabelmapEnvTerminationCriterion(min_reward_threshold=-0.1),
         reward_metric=LabelmapClusteringBasedReward(),
         project_actions_to="y",
+        apply_volume_transformation=True
     )
 
     experiment = (

scripts/armscan_dqn_sac_hl.py

@@ -35,7 +35,7 @@
         num_epochs=1,
         step_per_epoch=1000000,
         num_train_envs=-1,
-        num_test_envs=10,
+        num_test_envs=1,
         buffer_size=1000000,
         batch_size=256,
         step_per_collect=200,

scripts/armscan_sac_hl.py

@@ -33,8 +33,8 @@
     sampling_config = SamplingConfig(
         num_epochs=1,
         step_per_epoch=1000000,
-        num_train_envs=-1,
-        num_test_envs=10,
+        num_train_envs=40,
+        num_test_envs=1,
         buffer_size=1000000,
         batch_size=256,
         step_per_collect=200,