Reinforcement learning deployment in Xopt

.github/workflows/test-notebooks.yml

@@ -46,7 +46,7 @@ jobs:
           mpi: openmpi
 
       - name: Install dependencies
-        run: uv sync --extra dev
+        run: uv sync --extra dev --extra rl
 
       - name: Execute notebooks in parallel
         id: execute

.gitignore

@@ -2,6 +2,8 @@
 docs/examples/ga/nsga2/assets/yaml_runner_example/nsga2_output
 docs/examples/ga/nsga2/assets/yaml_runner_example/nsga2_from_checkpoint_output
 docs/examples/ga/nsga2/yaml_interface/assets/yaml_runner_example.zip
+docs/examples/rl/sac_pendulum.zip
+docs/examples/rl/inverted_pendulum_reward.png
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

docs/examples/rl/inverted_pendulum.ipynb

@@ -0,0 +1,295 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "646d8ad0",
+   "metadata": {},
+   "source": [
+    "# Deploying a trained RL policy in Xopt: inverted pendulum\n",
+    "\n",
+    "This notebook shows how to use a **reinforcement-learning policy trained externally**\n",
+    "(with [stable-baselines3](https://stable-baselines3.readthedocs.io/)) inside Xopt via\n",
+    "`RLGenerator`, which steps a gymnasium environment through `GymEvaluator`.\n",
+    "\n",
+    "Xopt does not train or update the policy here -- it is frozen at deployment time.\n",
+    "Training happens once, outside of Xopt (see `train_policy` below, or run\n",
+    "[`inverted_pendulum.py`](inverted_pendulum.py) directly), and the resulting policy is\n",
+    "simply used for inference by `RLGenerator.generate()`.\n",
+    "\n",
+    "Requires the optional `rl` extra: `pip install xopt[rl]`.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d287f697",
+   "metadata": {},
+   "source": [
+    "## Setup and imports\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "be07dc11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "import os\n",
+    "\n",
+    "import gymnasium as gym\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "from IPython.display import HTML\n",
+    "from matplotlib import animation\n",
+    "from stable_baselines3 import SAC\n",
+    "\n",
+    "from gest_api.vocs import VOCS\n",
+    "from xopt import Xopt\n",
+    "from xopt.evaluator import GymEvaluator\n",
+    "from xopt.generators.rl_generator import RLGenerator\n",
+    "from xopt.vocs import ContextualVariable\n",
+    "\n",
+    "MODEL_PATH = Path(\"sac_pendulum.zip\")\n",
+    "ACTION_NAMES = [\"torque\"]\n",
+    "OBSERVATION_NAMES = [\"cos_theta\", \"sin_theta\", \"theta_dot\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3399f268",
+   "metadata": {},
+   "source": [
+    "## Train (or load) the policy -- entirely outside of Xopt\n",
+    "\n",
+    "If `sac_pendulum.zip` already exists (e.g. from running `inverted_pendulum.py`), it is\n",
+    "loaded directly; otherwise a SAC policy is trained here for a modest number of timesteps.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0d80d8b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SMOKE_TEST = os.environ.get(\"SMOKE_TEST\")\n",
+    "\n",
+    "if MODEL_PATH.exists():\n",
+    "    policy = SAC.load(MODEL_PATH)\n",
+    "else:\n",
+    "    n_timesteps = 2 if SMOKE_TEST else 20_000\n",
+    "    train_env = gym.make(\"Pendulum-v1\")\n",
+    "    policy = SAC(\"MlpPolicy\", train_env, verbose=0)\n",
+    "    policy.learn(total_timesteps=n_timesteps)\n",
+    "    policy.save(MODEL_PATH)\n",
+    "    train_env.close()\n",
+    "\n",
+    "policy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b6f47ed3",
+   "metadata": {},
+   "source": [
+    "## Configure Xopt: VOCS, GymEvaluator, RLGenerator\n",
+    "\n",
+    "- `torque` is the action variable the policy controls.\n",
+    "- `cos_theta`, `sin_theta`, `theta_dot` are `ContextualVariable`s: contextual inputs to\n",
+    "  the policy that Xopt does not optimize over.\n",
+    "- `reward` is logged as the objective for bookkeeping/plotting only -- the frozen policy,\n",
+    "  not this objective, decides the next action.\n",
+    "\n",
+    "The environment is created with `render_mode=\"rgb_array\"` so we can grab frames for the\n",
+    "animation below. We also override the random reset with a non-ideal (but easily\n",
+    "recoverable) starting angle, so the rollout below first shows a quick recovery before\n",
+    "the later perturbation.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "edf5dcd3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vocs = VOCS(\n",
+    "    variables={\n",
+    "        \"torque\": [-2.0, 2.0],\n",
+    "        \"cos_theta\": ContextualVariable(),\n",
+    "        \"sin_theta\": ContextualVariable(),\n",
+    "        \"theta_dot\": ContextualVariable(),\n",
+    "    },\n",
+    "    objectives={\"reward\": \"MAXIMIZE\"},\n",
+    ")\n",
+    "\n",
+    "env = gym.make(\"Pendulum-v1\", render_mode=\"rgb_array\")\n",
+    "evaluator = GymEvaluator(\n",
+    "    env=env,\n",
+    "    action_space_names=ACTION_NAMES,\n",
+    "    observation_space_names=OBSERVATION_NAMES,\n",
+    "    max_workers=1,\n",
+    ")\n",
+    "\n",
+    "# start well off-vertical (but not a full hang-down swing-up) so recovery is quick\n",
+    "INITIAL_THETA = 2.0\n",
+    "INITIAL_THETA_DOT = 0.0\n",
+    "env.unwrapped.state = np.array([INITIAL_THETA, INITIAL_THETA_DOT])\n",
+    "evaluator._current_observation = np.array(\n",
+    "    [np.cos(INITIAL_THETA), np.sin(INITIAL_THETA), INITIAL_THETA_DOT]\n",
+    ")\n",
+    "\n",
+    "generator = RLGenerator(\n",
+    "    vocs=vocs,\n",
+    "    policy=policy,\n",
+    "    action_space_names=ACTION_NAMES,\n",
+    "    observation_space_names=OBSERVATION_NAMES,\n",
+    "    initial_observation=evaluator.current_observation,\n",
+    ")\n",
+    "\n",
+    "X = Xopt(generator=generator, evaluator=evaluator)\n",
+    "X"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "68a64504",
+   "metadata": {},
+   "source": [
+    "## Run the trained agent on the pendulum\n",
+    "\n",
+    "Step Xopt forward for one episode, capturing a render frame after every step for the\n",
+    "animation below. Halfway through, we manually kick the pendulum's angular velocity to\n",
+    "simulate an external disturbance and check that the frozen policy recovers.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8163417a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "N_STEPS = 200\n",
+    "PERTURBATION_STEP = N_STEPS // 2\n",
+    "PERTURBATION_THETA_DOT = 8.0\n",
+    "\n",
+    "frames = [env.render()]\n",
+    "for step in range(N_STEPS):\n",
+    "    X.step()\n",
+    "    if step == PERTURBATION_STEP:\n",
+    "        # simulate an external disturbance: kick the angular velocity mid-flight\n",
+    "        theta, theta_dot = env.unwrapped.state\n",
+    "        env.unwrapped.state = np.array([theta, theta_dot + PERTURBATION_THETA_DOT])\n",
+    "    frames.append(env.render())\n",
+    "\n",
+    "env.close()\n",
+    "X.data.tail()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7ec2c007",
+   "metadata": {},
+   "source": [
+    "## Visualize the state trajectory and reward\n",
+    "\n",
+    "The dashed line marks the mid-rollout disturbance; the policy was never trained on this\n",
+    "exact perturbation, so recovery afterward demonstrates its robustness.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2a299156",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = X.data.reset_index(drop=True)\n",
+    "theta = np.arctan2(data[\"sin_theta\"], data[\"cos_theta\"])\n",
+    "\n",
+    "fig, axes = plt.subplots(3, 1, figsize=(7, 7), sharex=True)\n",
+    "\n",
+    "axes[0].plot(theta)\n",
+    "axes[0].set_ylabel(r\"$\\theta$ (rad)\")\n",
+    "\n",
+    "axes[1].plot(data[\"theta_dot\"])\n",
+    "axes[1].set_ylabel(r\"$\\dot{\\theta}$ (rad/s)\")\n",
+    "\n",
+    "axes[2].plot(data[\"reward\"].cumsum())\n",
+    "axes[2].set_ylabel(\"cumulative reward\")\n",
+    "axes[2].set_xlabel(\"step\")\n",
+    "\n",
+    "for ax in axes:\n",
+    "    ax.axvline(PERTURBATION_STEP, color=\"red\", linestyle=\"--\", label=\"perturbation\")\n",
+    "axes[0].legend()\n",
+    "\n",
+    "fig.suptitle(\"RLGenerator rollout: frozen SAC policy recovering from a perturbation\")\n",
+    "fig.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "22214cd4",
+   "metadata": {},
+   "source": [
+    "## Animate the pendulum swing-up inline\n",
+    "\n",
+    "Renders the captured `rgb_array` frames as an inline JS animation (`to_jshtml`, no\n",
+    "`ffmpeg` required).\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4059503f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig_anim, ax_anim = plt.subplots(figsize=(4, 4))\n",
+    "ax_anim.axis(\"off\")\n",
+    "im = ax_anim.imshow(frames[0])\n",
+    "\n",
+    "\n",
+    "def _update(frame):\n",
+    "    im.set_data(frame)\n",
+    "    return (im,)\n",
+    "\n",
+    "\n",
+    "ani = animation.FuncAnimation(fig_anim, _update, frames=frames, interval=50, blit=True)\n",
+    "plt.close(fig_anim)\n",
+    "HTML(ani.to_jshtml())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3fea83d",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "xopt-dev",
+   "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.13.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}