Liquid time constant networks

CLAUDE.md

@@ -9,7 +9,7 @@ En cas de doute ou d'alternative de design : **poser la question d'abord**, prop
 options, attendre la réponse. Les corrections de bugs internes (fichiers `src/`) qui ne
 changent ni contrat ni comportement voulu restent autorisées.
 
-A neural network (SAC) trained to play a tank-arena game.
+A neural network (PPO) trained to play a tank-arena game.
 Jolt physics, offscreen Vulkan rendering (the agent's vision), NN via LibTorch.
 
 ## Architecture (CMake modules)
@@ -21,7 +21,9 @@ Dependency order: utils → model → view → core → agent / desktop
 - **arenai_view**   — Vulkan rendering; `offscreen_renderer` = offscreen render for the agent's vision
 - **arenai_core**   — `BaseTanksEnvironment`, enemy_handler, thread_pool (RL env loop)
 - **arenai_controller** — input handling
-- **arenai_agent**  — SAC: agents, networks, replay_buffer, reward_transforms + `main.cpp` (training executable)
+- **arenai_agent**  — agents (sac, ppo, ppo_liquid), networks, core (train env + spawn curriculum),
+  metrics + `main.cpp` (training executable). The reward is purely event-based
+  (hit/kill/received/death/win — no dense shaping).
 - **arenai_desktop** — the playable game executable. Its `src/` folders are hexagons of their own:
   `gui/` (RmlUi main menu — RmlUi types must never leak out of it, other code only includes
   `gui/menu.h`), `controller/`, `core/`. Menu assets live in `resources/menu/` + `resources/font/`.

arenai_agent/CMakeLists.txt

@@ -70,4 +70,6 @@ arenai_copy_runtime_dlls(arenai_agent_train TORCH)
 # Library
 add_library(arenai_agent SHARED ${ARENAI_AGENT_SRC})
 target_link_libraries(arenai_agent PRIVATE ${ARENAI_AGENT_LINK_LIBS} "${stb_SOURCE_DIR}")
+# nlohmann::json appears in the public factory.h contract
+target_link_libraries(arenai_agent PUBLIC nlohmann_json::nlohmann_json)
 target_include_directories(arenai_agent PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/include")

arenai_agent/include/arenai_agent/factory.h

@@ -5,64 +5,52 @@
 #ifndef ARENAI_AGENT_HOST_FACTORY_H
 #define ARENAI_AGENT_HOST_FACTORY_H
 
-#include <format>
-#include <map>
 #include <memory>
 #include <string>
 
+#include <nlohmann/json.hpp>
+
 #include "./agent.h"
 
 namespace arenai::agent {
 
+    enum AgentAlgorithm { PPO, PPO_LIQUID };
+
     class AgentFactory {
     public:
-        virtual ~AgentFactory() = default;
-
-        explicit AgentFactory(const std::map<std::string, std::string> &arguments);
+        // config: the content of a training run's config.json — the network
+        // hyper-parameters come from its "agent" section, the vision size and
+        // the control frequency from its "environment" section
+        explicit AgentFactory(const nlohmann::json &config);
 
         std::shared_ptr<AbstractAgent> get_agent(
-            const int &vision_height, const int &vision_width, const int &nb_sensors,
-            const int &nb_continuous_actions, const int &nb_discrete_actions);
-
-    protected:
-        template<typename T>
-        T get_value(const std::string &argument_name, T default_value) {
-            if (!arguments.contains(argument_name)) return default_value;
-
-            const std::string value_as_string = arguments[argument_name];
-            std::stringstream ss(value_as_string);
-            T value;
-            ss >> value;
+            AgentAlgorithm algorithm, const int &nb_sensors, const int &nb_continuous_actions,
+            const int &nb_discrete_actions, bool cuda);
 
-            if (ss.fail() || !ss.eof())
-                throw std::runtime_error(std::format(
-                    R"(Wrong value for "{}" : "{}", example : "{}")", argument_name,
-                    value_as_string, default_value));
-
-            arguments.erase(arguments.find(argument_name));
-
-            return value;
-        }
+        int get_vision_height() const;
+        int get_vision_width() const;
+        float get_wanted_frequency() const;
 
+    private:
+        // every hyper-parameter is required: a config.json that misses one
+        // does not describe the run it claims to (at() throws on a missing key)
         template<typename T>
-        T get_value(
-            const std::string &argument_name, const std::function<T(std::string)> &parse_fn,
-            T default_value) {
-            if (!arguments.contains(argument_name)) return default_value;
-
-            const std::string value_as_string = arguments[argument_name];
-
-            arguments.erase(arguments.find(argument_name));
-
-            return parse_fn(value_as_string);
+        T get_value(const std::string &argument_name) {
+            return agent_arguments.at(argument_name).get<T>();
         }
 
-        virtual std::shared_ptr<AbstractAgent> get_agent_impl(
-            const int &vision_height, const int &vision_width, const int &nb_sensors,
-            const int &nb_continuous_actions, const int &nb_discrete_action) = 0;
+        std::shared_ptr<AbstractAgent> create_ppo_agent(
+            const int &nb_sensors, const int &nb_continuous_actions, const int &nb_discrete_action,
+            bool cuda);
 
-    private:
-        std::map<std::string, std::string> arguments;
+        std::shared_ptr<AbstractAgent> create_liquid_ppo_agent(
+            const int &nb_sensors, const int &nb_continuous_actions, const int &nb_discrete_action,
+            bool cuda);
+
+        nlohmann::json agent_arguments;
+        int vision_height;
+        int vision_width;
+        float wanted_frequency;
     };
 
 }// namespace arenai::agent

arenai_agent/python/.gitignore

@@ -1 +1,2 @@
 __pycache__/
+.ipynb_checkpoints/

arenai_agent/python/notebooks/convert_liquid_cell_fusion.ipynb

@@ -0,0 +1,206 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "intro",
+   "metadata": {},
+   "source": [
+    "# Liquid checkpoint conversion — `LiquidRecurrent` / `LiquidCell` fusion\n",
+    "\n",
+    "The classes `LiquidRecurrent` and `LiquidCell` were merged into a single `LiquidCell`\n",
+    "(`arenai_agent/src/networks/recurrent/liquid_cell.h`). The intermediate `cell` submodule\n",
+    "disappeared, so every parameter saved under `liquid.cell.*` now lives under `liquid.*`:\n",
+    "\n",
+    "| old key | new key |\n",
+    "|---|---|\n",
+    "| `liquid.cell.a` | `liquid.a` |\n",
+    "| `liquid.cell.raw_tau` | `liquid.raw_tau` |\n",
+    "| `liquid.cell.f.*` | `liquid.f.*` |\n",
+    "| `liquid.to_output.*` | unchanged |\n",
+    "\n",
+    "This notebook reads an `actor.pt` saved **before** the fusion (via\n",
+    "`torch::serialize::OutputArchive`), remaps the keys and writes a checkpoint loadable by the\n",
+    "fused C++ code (`torch::serialize::InputArchive`). The same mapping applies to `critic.pt`.\n",
+    "\n",
+    "Notes:\n",
+    "- the parameter **order** is unchanged by the fusion, so `actor_optim.pt` / `critic_optim.pt`\n",
+    "  need no conversion;\n",
+    "- requires `torch` (`pip install torch`, the CPU wheel is enough — it is not part of the\n",
+    "  `arenai_agent/python` requirements)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "paths",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import re\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "\n",
+    "OLD_CHECKPOINT = Path(\n",
+    "    \"/home/samuel/CLionProjects/ArenAI/outputs/train_400_ppo_liquid_beta_silu_truncation-no-penalty/save_97/actor.pt\"\n",
+    ")\n",
+    "# written under a `converted/` sibling folder, same file name\n",
+    "NEW_CHECKPOINT = OLD_CHECKPOINT.parent / \"converted\" / OLD_CHECKPOINT.name"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "load-md",
+   "metadata": {},
+   "source": [
+    "## Read the old state dict\n",
+    "\n",
+    "A module saved by LibTorch's `OutputArchive` is a TorchScript archive: `torch.jit.load`\n",
+    "gives back the full module hierarchy with its named parameters."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "load",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "old_module = torch.jit.load(str(OLD_CHECKPOINT), map_location=\"cpu\")\n",
+    "\n",
+    "old_params = {name: p.detach().clone() for name, p in old_module.named_parameters()}\n",
+    "old_buffers = dict(old_module.named_buffers())\n",
+    "assert not old_buffers, f\"unexpected buffers, extend the notebook: {sorted(old_buffers)}\"\n",
+    "\n",
+    "for name, p in old_params.items():\n",
+    "    print(f\"{name:45s} {tuple(p.shape)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "convert-md",
+   "metadata": {},
+   "source": [
+    "## Remap and rebuild\n",
+    "\n",
+    "The C++ `Module::load` walks the archive by **module hierarchy**, and requires every\n",
+    "serialized submodule to exist — including parameterless ones (SiLU, Sigmoid, LayerNorm\n",
+    "positions in the `Sequential`s...). So the converted archive must replicate the full module\n",
+    "tree, not just the parameter paths: the old `liquid.cell` node is dropped and its children\n",
+    "are re-attached directly to `liquid`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "convert",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def convert(name: str) -> str:\n",
+    "    return re.sub(r\"^liquid\\.cell\\.\", \"liquid.\", name)\n",
+    "\n",
+    "\n",
+    "class Node(nn.Module):\n",
+    "    def forward(self) -> None:\n",
+    "        return None\n",
+    "\n",
+    "\n",
+    "def ensure_path(root: Node, dotted: str) -> Node:\n",
+    "    module = root\n",
+    "    for part in dotted.split(\".\"):\n",
+    "        if part not in module._modules:\n",
+    "            module.add_module(part, Node())\n",
+    "        module = module._modules[part]\n",
+    "    return module\n",
+    "\n",
+    "\n",
+    "root = Node()\n",
+    "for name, _ in old_module.named_modules():\n",
+    "    if name and name != \"liquid.cell\":\n",
+    "        ensure_path(root, convert(name))\n",
+    "for name, tensor in old_params.items():\n",
+    "    *path, leaf = convert(name).split(\".\")\n",
+    "    ensure_path(root, \".\".join(path)).register_parameter(leaf, nn.Parameter(tensor))\n",
+    "\n",
+    "NEW_CHECKPOINT.parent.mkdir(parents=True, exist_ok=True)\n",
+    "torch.jit.save(torch.jit.script(root), str(NEW_CHECKPOINT))\n",
+    "print(f\"written: {NEW_CHECKPOINT}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "verify-md",
+   "metadata": {},
+   "source": [
+    "## Verify the round trip"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "verify",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "reloaded = torch.jit.load(str(NEW_CHECKPOINT), map_location=\"cpu\")\n",
+    "\n",
+    "new_params = dict(reloaded.named_parameters())\n",
+    "expected = {convert(name): t for name, t in old_params.items()}\n",
+    "\n",
+    "assert set(new_params) == set(expected), set(new_params) ^ set(expected)\n",
+    "for name, t in expected.items():\n",
+    "    assert torch.equal(new_params[name], t), name\n",
+    "\n",
+    "# structure the C++ InputArchive will walk\n",
+    "liquid = reloaded.liquid\n",
+    "for attr in (\"a\", \"raw_tau\", \"f\", \"to_output\"):\n",
+    "    assert hasattr(liquid, attr), attr\n",
+    "assert not hasattr(liquid, \"cell\")\n",
+    "\n",
+    "print(f\"OK — {len(new_params)} parameters remapped\")\n",
+    "for name in sorted(new_params):\n",
+    "    if name.startswith(\"liquid\"):\n",
+    "        print(f\"  {name:35s} {tuple(new_params[name].shape)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "98c45cbd-59e2-4ae7-8918-54e749510121",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0cf3755b-aebe-45f6-b8ee-c8caad225b5c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.14.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

arenai_agent/src/agents/agent_cli.cpp

@@ -9,8 +9,8 @@
 #include "../utils/cli_fields.h"
 #include "./ppo/ppo_factory.h"
 #include "./ppo/ppo_hyperparams.h"
-#include "./sac/sac_factory.h"
-#include "./sac/sac_hyperparams.h"
+#include "./ppo_liquid/liquid_ppo_factory.h"
+#include "./ppo_liquid/liquid_ppo_hyperparams.h"
 
 namespace arenai::agent {
 
@@ -39,10 +39,10 @@ namespace arenai::agent {
     std::vector<AgentCli> make_agent_clis() {
         std::vector<AgentCli> algorithms;
 
-        algorithms.push_back(
-            make_agent_cli<SacHyperParams, SacTorchAgentFactory>("sac", sac_cli_fields()));
         algorithms.push_back(
             make_agent_cli<PpoHyperParams, PpoTorchAgentFactory>("ppo", ppo_cli_fields()));
+        algorithms.push_back(make_agent_cli<LiquidPpoHyperParams, LiquidPpoTorchAgentFactory>(
+            "ppo_liquid", liquid_ppo_cli_fields()));
 
         return algorithms;
     }