diff --git a/.github/workflows/test-notebooks.yml b/.github/workflows/test-notebooks.yml index 64ddf12b1..296f39d8a 100644 --- a/.github/workflows/test-notebooks.yml +++ b/.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 diff --git a/.gitignore b/.gitignore index 78ebf2da5..ce5baf49f 100644 --- a/.gitignore +++ b/.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__/ diff --git a/docs/examples/rl/inverted_pendulum.ipynb b/docs/examples/rl/inverted_pendulum.ipynb new file mode 100644 index 000000000..7d693d8f4 --- /dev/null +++ b/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 +} diff --git a/docs/examples/rl/inverted_pendulum.py b/docs/examples/rl/inverted_pendulum.py new file mode 100644 index 000000000..870cb2c89 --- /dev/null +++ b/docs/examples/rl/inverted_pendulum.py @@ -0,0 +1,92 @@ +""" +Reinforcement-learning example: swing up an inverted pendulum with Xopt. + +This script has two independent parts: + +1. ``train_policy`` trains a SAC policy on gymnasium's ``Pendulum-v1`` with + stable-baselines3, entirely outside of Xopt (the policy is just saved to disk). +2. ``run_with_xopt`` loads that frozen policy and deploys it through + ``RLGenerator`` + ``GymEvaluator`` to run one Xopt-driven rollout. + +Run with: ``python docs/examples/rl/inverted_pendulum.py`` +Requires the optional ``rl`` extra: ``pip install xopt[rl]``. +""" + +from pathlib import Path + +import gymnasium as gym +import matplotlib.pyplot as plt +from stable_baselines3 import SAC + +from gest_api.vocs import VOCS +from xopt import Xopt +from xopt.evaluator import GymEvaluator +from xopt.generators.rl_generator import RLGenerator +from xopt.vocs import ContextualVariable + +MODEL_PATH = Path(__file__).parent / "sac_pendulum.zip" +ACTION_NAMES = ["torque"] +OBSERVATION_NAMES = ["cos_theta", "sin_theta", "theta_dot"] + + +def train_policy() -> SAC: + """Train (or load a previously cached) SAC policy on Pendulum-v1.""" + if MODEL_PATH.exists(): + return SAC.load(MODEL_PATH) + + env = gym.make("Pendulum-v1") + model = SAC("MlpPolicy", env, verbose=0) + model.learn(total_timesteps=20_000) + model.save(MODEL_PATH) + env.close() + return model + + +def run_with_xopt(policy: SAC, n_steps: int = 200): + """Deploy the frozen policy inside Xopt via RLGenerator + GymEvaluator.""" + vocs = VOCS( + variables={ + "torque": [-2.0, 2.0], + "cos_theta": ContextualVariable(), + "sin_theta": ContextualVariable(), + "theta_dot": ContextualVariable(), + }, + objectives={"reward": "MAXIMIZE"}, + ) + + env = gym.make("Pendulum-v1") + evaluator = GymEvaluator( + env=env, + action_space_names=ACTION_NAMES, + observation_space_names=OBSERVATION_NAMES, + max_workers=1, + ) + generator = RLGenerator( + vocs=vocs, + policy=policy, + action_space_names=ACTION_NAMES, + observation_space_names=OBSERVATION_NAMES, + initial_observation=evaluator.current_observation, + ) + + X = Xopt(generator=generator, evaluator=evaluator) + for _ in range(n_steps): + X.step() + + env.close() + return X.data + + +if __name__ == "__main__": + trained_policy = train_policy() + data = run_with_xopt(trained_policy) + + print(f"ran {len(data)} steps, cumulative reward = {data['reward'].sum():.2f}") + + plt.plot(data["reward"].cumsum()) + plt.xlabel("step") + plt.ylabel("cumulative reward") + plt.title("Inverted pendulum: RLGenerator rollout with a frozen SAC policy") + plot_path = Path(__file__).parent / "inverted_pendulum_reward.png" + plt.savefig(plot_path) + print(f"saved plot to {plot_path}") diff --git a/pyproject.toml b/pyproject.toml index 0fbb7b02e..7b0514d9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "tqdm", "orjson", "matplotlib", - "gest-api>=0.2" + "gest-api" ] description = "Flexible optimization of arbitrary problems in Python." dynamic = [ "version" ] @@ -66,6 +66,11 @@ doc = [ "mkdocstrings", "mkdocstrings-python", ] +rl = [ + "gymnasium", + "stable-baselines3", + "pygame", +] [project.urls] Homepage = "https://github.com/xopt-org/xopt" diff --git a/xopt/evaluator.py b/xopt/evaluator.py index ed7768237..ba169d149 100644 --- a/xopt/evaluator.py +++ b/xopt/evaluator.py @@ -2,12 +2,15 @@ from concurrent.futures import Executor, Future, ProcessPoolExecutor from enum import Enum from threading import Lock -from typing import Callable, Dict, List, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union import numpy as np import pandas as pd from pandas import DataFrame -from pydantic import ConfigDict, Field, model_validator +from pydantic import ConfigDict, Field, PrivateAttr, model_validator + +if TYPE_CHECKING: + import gymnasium as gym from xopt.errors import XoptError from xopt.pydantic import NormalExecutor, XoptBaseModel @@ -437,3 +440,149 @@ def shutdown(self, wait: bool = True): """ with self._shutdownLock: self._shutdown = True + + +try: + import gymnasium as gym # noqa: F811 + + _HAS_GYMNASIUM = True +except ModuleNotFoundError: + _HAS_GYMNASIUM = False + +if not _HAS_GYMNASIUM: + logger.debug("gymnasium not installed, GymEvaluator is not available") +else: + + class GymEvaluator(Evaluator): + """ + Evaluator for OpenAI Gym environments. + + Parameters + ---------- + env : gym.Env + The Gym environment. + action_space_names : List[str] + Names of action space dimensions. + observation_space_names : List[str] + Names of observation space dimensions. + """ + + env: gym.Env + action_space_names: List[str] + observation_space_names: List[str] + action_history: List[np.ndarray] = [] + action_mode: str + + # state the environment is in right now, i.e. the state the *next* action + # will be applied to. Tracked internally since env.step() only returns the + # state *after* the action, and env.reset()'s return value would otherwise + # be discarded. + _current_observation: Optional[np.ndarray] = PrivateAttr(default=None) + + def __init__( + self, + env: gym.Env, + action_space_names: List[str], + observation_space_names: List[str], + **kwargs, + ): + """ + Initialize the GymEvaluator. + + Parameters + ---------- + env : gym.Env + The name of the environment. + action_space_names : List[str] + Names of action space dimensions. + observation_space_names : List[str] + Names of observation space dimensions. + """ + + if hasattr(env, "action_mode"): + action_mode = env.action_mode + else: + action_mode = "delta" + + logger.debug(f"Using action mode: {action_mode}") + + function = self._evaluate_function + + super().__init__( + env=env, + action_space_names=action_space_names, + observation_space_names=observation_space_names, + function=function, + action_mode=action_mode, + **kwargs, + ) + self.reset() + + def reset(self): + """ + Reset the environment and clear the action history. + """ + observation, _ = self.env.reset() + self._current_observation = observation + self.action_history = [] + + @property + def current_observation(self) -> Dict[str, float]: + """Current observation dict, i.e. the state the next action will be applied to.""" + return { + name: float(self._current_observation[i]) + for i, name in enumerate(self.observation_space_names) + } + + def _evaluate_function(self, x: dict) -> dict: + """ + Evaluate the given input using the environment. + + Parameters + ---------- + x : dict + The input dictionary to evaluate. + + Returns + ------- + dict + The evaluation results. + """ + xopt_action = np.array([x[name] for name in self.action_space_names]) + + # if self.action_mode == "delta": + # action = target_state - self.current_action_state + # else: + # action = target_state + action = xopt_action + + # snapshot the state the action is actually applied to -- env.step() + # computes reward from this state, not from the state it returns, so + # pairing the action with the returned observation would associate it + # with the wrong (resulting, not originating) state. + state = self._current_observation + + observation, reward, terminated, truncated, info = self.env.step(action) + self.action_history.append(action) + self._current_observation = observation + + observations = { + name: float(state[i]) + for i, name in enumerate(self.observation_space_names) + } + next_observations = { + f"next_{name}": float(observation[i]) + for i, name in enumerate(self.observation_space_names) + } + + # TODO: Multi-objective support + return ( + { + "reward": reward, + "terminated": terminated, + "truncated": truncated, + "info": info, + } + | observations + | next_observations + ) diff --git a/xopt/generators/__init__.py b/xopt/generators/__init__.py index 776ed8d32..f4d1794de 100644 --- a/xopt/generators/__init__.py +++ b/xopt/generators/__init__.py @@ -28,6 +28,7 @@ "ga": {"cnsga", "nsga2"}, "es": {"extremum_seeking"}, "rcds": {"rcds"}, + "rl": {"rl_policy"}, } @@ -125,6 +126,11 @@ def get_generator_dynamic(name: str) -> type[Generator]: generators[name] = RCDSGenerator return RCDSGenerator + elif name in all_generator_names["rl"]: + from xopt.generators.rl_generator import RLGenerator + + generators[name] = RLGenerator + return RLGenerator raise KeyError diff --git a/xopt/generators/rl_generator.py b/xopt/generators/rl_generator.py new file mode 100644 index 000000000..95dfa6f8e --- /dev/null +++ b/xopt/generators/rl_generator.py @@ -0,0 +1,113 @@ +import logging +from typing import Any, ClassVar, Dict, List + +import numpy as np +from pydantic import ConfigDict, model_validator + +from xopt.errors import GeneratorError, VOCSError +from xopt.generator import Generator +from xopt.vocs import ContextualVariable + +from gest_api.vocs import ContinuousVariable + +logger = logging.getLogger(__name__) + + +class RLGenerator(Generator): + """ + Generator that deploys a reinforcement-learning policy trained externally + (outside Xopt) against a stateful evaluator such as `GymEvaluator`. + + The policy must expose a stable-baselines3 style + ``predict(observation, deterministic) -> (action, state)`` method. Xopt does + not train or update the policy; it is only used for inference. + + Parameters + ---------- + policy : Any + Trained policy object with a ``predict(observation, deterministic)`` method. + action_space_names : List[str] + VOCS variable names the policy's action maps to, in policy output order. + observation_space_names : List[str] + VOCS variable names the policy's observation is built from, in policy input order. + initial_observation : Dict[str, float] + Observation used to select the very first action, before any data exists. + deterministic : bool, default=True + Whether to sample the policy deterministically. + """ + + name: ClassVar[str] = "rl_policy" + supports_batch_generation: bool = False + supports_single_objective: bool = True + supports_multi_objective: bool = False + supports_constraints: bool = False + supports_discrete_variables: bool = False + supports_contextual_variables: bool = True + + policy: Any + action_space_names: List[str] + observation_space_names: List[str] + initial_observation: Dict[str, float] + deterministic: bool = True + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @model_validator(mode="after") + def _validate_rl_names(self): + for name in self.action_space_names: + if name not in self.vocs.variables: + raise VOCSError(f"action variable `{name}` not found in vocs.variables") + if isinstance(self.vocs.variables[name], ContextualVariable): + raise VOCSError( + f"action variable `{name}` cannot be a ContextualVariable" + ) + if not isinstance(self.vocs.variables[name], ContinuousVariable): + raise VOCSError( + f"action variable `{name}` must be a ContinuousVariable" + ) + + for name in self.observation_space_names: + if name not in self.vocs.variables: + raise VOCSError( + f"observation variable `{name}` not found in vocs.variables" + ) + if not isinstance(self.vocs.variables[name], ContextualVariable): + raise VOCSError( + f"observation variable `{name}` must be a ContextualVariable" + ) + + missing = set(self.observation_space_names) - set(self.initial_observation) + if missing: + raise VOCSError(f"initial_observation is missing entries for {missing}") + + return self + + def _current_observation(self) -> Dict[str, float]: + """Latest known observation, from the last evaluated row or the initial fallback.""" + if self.data is None or len(self.data) == 0: + return self.initial_observation + + last_row = self.data.iloc[-1] + return { + name: float(last_row[f"next_{name}"]) + for name in self.observation_space_names + } + + def generate(self, n_candidates: int) -> List[Dict[str, float]]: + if n_candidates != 1: + raise GeneratorError( + "RLGenerator only supports generating one candidate at a time" + ) + + observation = self._current_observation() + obs_array = np.array( + [observation[name] for name in self.observation_space_names] + ) + + action, _state = self.policy.predict( + obs_array, deterministic=self.deterministic + ) + + return [ + {name: float(action[i]) for i, name in enumerate(self.action_space_names)} + ] diff --git a/xopt/tests/generators/test_rl_generator.py b/xopt/tests/generators/test_rl_generator.py new file mode 100644 index 000000000..0b7fff3e6 --- /dev/null +++ b/xopt/tests/generators/test_rl_generator.py @@ -0,0 +1,108 @@ +import numpy as np +import pytest + +from xopt import Xopt +from xopt.errors import GeneratorError, VOCSError +from xopt.generators.rl_generator import RLGenerator +from xopt.vocs import ContextualVariable + +from gest_api.vocs import VOCS + + +class StubPolicy: + """Fixed-action stand-in for a trained stable-baselines3 policy.""" + + def predict(self, observation, deterministic=True): + return np.array([0.0]), None + + +def _build_vocs(): + return VOCS( + variables={ + "torque": [-2.0, 2.0], + "cos_theta": ContextualVariable(), + "sin_theta": ContextualVariable(), + "theta_dot": ContextualVariable(), + }, + objectives={"reward": "MAXIMIZE"}, + ) + + +class TestRLGenerator: + def test_generate_uses_initial_observation(self): + vocs = _build_vocs() + obs_names = ["cos_theta", "sin_theta", "theta_dot"] + generator = RLGenerator( + vocs=vocs, + policy=StubPolicy(), + action_space_names=["torque"], + observation_space_names=obs_names, + initial_observation={"cos_theta": 1.0, "sin_theta": 0.0, "theta_dot": 0.0}, + ) + + candidates = generator.generate(1) + assert candidates == [{"torque": 0.0}] + + def test_generate_rejects_batch(self): + vocs = _build_vocs() + generator = RLGenerator( + vocs=vocs, + policy=StubPolicy(), + action_space_names=["torque"], + observation_space_names=["cos_theta", "sin_theta", "theta_dot"], + initial_observation={"cos_theta": 1.0, "sin_theta": 0.0, "theta_dot": 0.0}, + ) + with pytest.raises(GeneratorError): + generator.generate(2) + + def test_action_name_must_be_continuous(self): + vocs = _build_vocs() + with pytest.raises(VOCSError): + RLGenerator( + vocs=vocs, + policy=StubPolicy(), + action_space_names=["cos_theta"], + observation_space_names=["sin_theta", "theta_dot"], + initial_observation={"sin_theta": 0.0, "theta_dot": 0.0}, + ) + + def test_observation_name_must_be_contextual_variable(self): + vocs = _build_vocs() + with pytest.raises(VOCSError): + RLGenerator( + vocs=vocs, + policy=StubPolicy(), + action_space_names=["torque"], + observation_space_names=["torque"], + initial_observation={"torque": 0.0}, + ) + + def test_run_with_gym_evaluator(self): + gym = pytest.importorskip("gymnasium") + from xopt.evaluator import GymEvaluator + + vocs = _build_vocs() + obs_names = ["cos_theta", "sin_theta", "theta_dot"] + env = gym.make("Pendulum-v1") + evaluator = GymEvaluator( + env=env, + action_space_names=["torque"], + observation_space_names=obs_names, + max_workers=1, + ) + generator = RLGenerator( + vocs=vocs, + policy=StubPolicy(), + action_space_names=["torque"], + observation_space_names=obs_names, + initial_observation=evaluator.current_observation, + ) + + X = Xopt(generator=generator, evaluator=evaluator) + for _ in range(3): + X.step() + + assert len(X.data) == 3 + for name in obs_names: + assert f"next_{name}" in X.data.columns + assert all(X.data["torque"] == 0.0) diff --git a/xopt/tests/test_evaluator.py b/xopt/tests/test_evaluator.py index 8029c30f4..f3158d492 100644 --- a/xopt/tests/test_evaluator.py +++ b/xopt/tests/test_evaluator.py @@ -192,3 +192,35 @@ def test_evaluate_data_vectorized(self): assert result.shape[0] == 5 assert "f" in result.columns assert all(result["f"] == candidates["x1"] ** 2 + candidates["x2"] ** 2) + + +class TestGymEvaluator: + def test_current_observation_and_step(self): + gym = pytest.importorskip("gymnasium") + from xopt.evaluator import GymEvaluator + + env = gym.make("Pendulum-v1") + obs_names = ["cos_theta", "sin_theta", "theta_dot"] + evaluator = GymEvaluator( + env=env, + action_space_names=["torque"], + observation_space_names=obs_names, + ) + + # current_observation should match the reset observation + env.reset(seed=0) + evaluator.reset() + assert evaluator.current_observation == { + name: pytest.approx(float(evaluator._current_observation[i])) + for i, name in enumerate(obs_names) + } + + result = evaluator.evaluate({"torque": 0.0}) + assert set(obs_names).issubset(result) + assert {f"next_{n}" for n in obs_names}.issubset(result) + assert "reward" in result + + # current_observation should now reflect the post-step state + assert evaluator.current_observation == { + name: pytest.approx(result[f"next_{name}"]) for name in obs_names + } diff --git a/xopt/tests/test_generator.py b/xopt/tests/test_generator.py index 98cb237c4..322a19a78 100644 --- a/xopt/tests/test_generator.py +++ b/xopt/tests/test_generator.py @@ -16,6 +16,13 @@ from gest_api.vocs import VOCS +class _StubPolicy: + """Fixed-action stand-in for a trained RL policy, used in generator serialization tests.""" + + def predict(self, observation, deterministic=True): + return [0.0], None + + class PatchGenerator(Generator): """ Test generator class for testing purposes. @@ -115,6 +122,18 @@ def test_serialization_loading(self, name): ).model_dump() json.dumps(gen_config) + gen_class(vocs=test_vocs, **gen_config) + elif name in ["rl_policy"]: + # policy is an externally-trained runtime object, not JSON-serializable by design + test_vocs = VOCS( + variables={"x1": [0, 1], "obs": ContextualVariable()}, + objectives={"y1": "MINIMIZE"}, + ) + gen_config["policy"] = _StubPolicy() + gen_config["action_space_names"] = ["x1"] + gen_config["observation_space_names"] = ["obs"] + gen_config["initial_observation"] = {"obs": 0.0} + gen_class(vocs=test_vocs, **gen_config) else: test_vocs = deepcopy(TEST_VOCS_BASE)