From 8d752a16a52b5848d999b363b2c192cedba53759 Mon Sep 17 00:00:00 2001 From: Vishwa GW <136871744+vishwagw@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:29:43 +0530 Subject: [PATCH 1/3] Add DQN training example notebook --- examples/dqn_training.ipynb | 530 ++++++++++++++++++++++++++++++++++++ 1 file changed, 530 insertions(+) create mode 100644 examples/dqn_training.ipynb diff --git a/examples/dqn_training.ipynb b/examples/dqn_training.ipynb new file mode 100644 index 00000000..9b2429ee --- /dev/null +++ b/examples/dqn_training.ipynb @@ -0,0 +1,530 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Training a DQN Agent\n", + "The other training examples use [RLlib](rllib_training.ipynb) with PPO. BSK-RL environments\n", + "are ordinary [Gymnasium](https://gymnasium.farama.org/) environments, so they also work with\n", + "hand-written training loops. This example trains a\n", + "[Deep Q-Network](https://www.nature.com/articles/nature14236) (DQN) in plain PyTorch, with no\n", + "RL library.\n", + "\n", + "DQN fits BSK-RL well for two reasons:\n", + "\n", + "* **Discrete actions.** Most satellite action specs (scan, charge, downlink, desaturate,\n", + " image target *n*) give a ``Discrete`` action space, which is what DQN needs.\n", + "* **Off-policy learning.** Simulator steps are expensive. A replay buffer lets every\n", + " transition be reused for many gradient updates.\n", + "\n", + "## Discounting in a semi-MDP\n", + "BSK-RL steps do not all last the same amount of time: a scan here lasts 180 seconds and a\n", + "downlink lasts 60. As described in the [time-discounted GAE](time_discounted_gae.ipynb)\n", + "example, the discount should depend on the elapsed time, not the number of steps. For\n", + "one-step Q-learning, the target becomes\n", + "\n", + "\\begin{equation}\n", + "y = r + \\gamma^{\\Delta t} \\, (1 - \\text{terminated}) \\, \\max_{a'} Q_\\text{target}(s', a')\n", + "\\end{equation}\n", + "\n", + "where $\\Delta t$ is the step duration in seconds, read from ``info[\"d_ts\"]``, and $\\gamma$ is a\n", + "*per-second* discount factor. This is the same convention as\n", + "``TimeDiscountedGAEPPOTorchLearner``, with reward credited at the end of the step.\n", + "\n", + "## Load Modules" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import random\n", + "from typing import ClassVar\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from Basilisk.architecture import bskLogging\n", + "from torch import nn\n", + "\n", + "from bsk_rl import SatelliteTasking, act, data, obs, sats, scene\n", + "from bsk_rl.sim import dyn, fsw\n", + "\n", + "bskLogging.setDefaultLogLevel(bskLogging.BSK_WARNING)\n", + "\n", + "SEED = 0\n", + "random.seed(SEED)\n", + "np.random.seed(SEED)\n", + "torch.manual_seed(SEED)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Define the Environment\n", + "The environment is the nadir-scanning and downlink mission from the\n", + "[RLlib example](rllib_training.ipynb), so the two examples can be compared directly. The\n", + "satellite collects data while keeping its data buffer, battery, and reaction wheel speeds\n", + "within limits. If any of these limits is violated, the satellite fails, the episode ends,\n", + "and a penalty is applied." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "class ScanningDownlinkDynModel(\n", + " dyn.ContinuousImagingDynModel, dyn.GroundStationDynModel\n", + "):\n", + " # Define some custom properties to be accessed in the state\n", + " @property\n", + " def instrument_pointing_error(self) -> float:\n", + " r_BN_P_unit = self.r_BN_P / np.linalg.norm(self.r_BN_P)\n", + " c_hat_P = self.satellite.fsw.c_hat_P\n", + " return np.arccos(np.dot(-r_BN_P_unit, c_hat_P))\n", + "\n", + " @property\n", + " def solar_pointing_error(self) -> float:\n", + " a = (\n", + " self.world.gravFactory.spiceObject.planetStateOutMsgs[self.world.sun_index]\n", + " .read()\n", + " .PositionVector\n", + " )\n", + " a_hat_N = a / np.linalg.norm(a)\n", + " nHat_B = self.satellite.sat_args[\"nHat_B\"]\n", + " NB = np.transpose(self.BN)\n", + " nHat_N = NB @ nHat_B\n", + " return np.arccos(np.dot(nHat_N, a_hat_N))\n", + "\n", + "\n", + "class ScanningSatellite(sats.AccessSatellite):\n", + " observation_spec: ClassVar[list[obs.Observation]] = [\n", + " obs.SatProperties(\n", + " dict(prop=\"storage_level_fraction\"),\n", + " dict(prop=\"battery_charge_fraction\"),\n", + " dict(prop=\"wheel_speeds_fraction\"),\n", + " dict(prop=\"instrument_pointing_error\", norm=np.pi),\n", + " dict(prop=\"solar_pointing_error\", norm=np.pi),\n", + " ),\n", + " obs.OpportunityProperties(\n", + " dict(prop=\"opportunity_open\", norm=5700),\n", + " dict(prop=\"opportunity_close\", norm=5700),\n", + " type=\"ground_station\",\n", + " n_ahead_observe=1,\n", + " ),\n", + " obs.Eclipse(norm=5700),\n", + " obs.Time(),\n", + " ]\n", + " action_spec: ClassVar[list[act.Action]] = [\n", + " act.Scan(duration=180.0),\n", + " act.Charge(duration=120.0),\n", + " act.Downlink(duration=60.0),\n", + " act.Desat(duration=60.0),\n", + " ]\n", + " dyn_type = ScanningDownlinkDynModel\n", + " fsw_type = fsw.ContinuousImagingFSWModel" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The satellite parameters and environment arguments match the RLlib example. Because\n", + "``terminate_on_time_limit=True``, reaching the end of the mission is a true terminal\n", + "state, so the agent does not bootstrap past it." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "sat = ScanningSatellite(\n", + " \"Scanner_1\",\n", + " sat_args=dict(\n", + " # Data\n", + " dataStorageCapacity=5000 * 8e6, # bits\n", + " storageInit=lambda: np.random.uniform(0.0, 0.8) * 5000 * 8e6,\n", + " instrumentBaudRate=0.5 * 8e6,\n", + " transmitterBaudRate=-50 * 8e6,\n", + " # Power\n", + " batteryStorageCapacity=200 * 3600, # W*s\n", + " storedCharge_Init=lambda: np.random.uniform(0.3, 1.0) * 200 * 3600,\n", + " basePowerDraw=-10.0, # W\n", + " instrumentPowerDraw=-30.0, # W\n", + " transmitterPowerDraw=-25.0, # W\n", + " thrusterPowerDraw=-80.0, # W\n", + " panelArea=0.25,\n", + " # Attitude\n", + " imageAttErrorRequirement=0.1,\n", + " imageRateErrorRequirement=0.1,\n", + " disturbance_vector=lambda: np.random.normal(scale=0.0001, size=3), # N*m\n", + " maxWheelSpeed=6000.0, # RPM\n", + " wheelSpeeds=lambda: np.random.uniform(-3000, 3000, 3),\n", + " desatAttitude=\"nadir\",\n", + " ),\n", + ")\n", + "\n", + "duration = 5 * 5700.0 # About 5 orbits\n", + "env = SatelliteTasking(\n", + " satellite=sat,\n", + " scenario=scene.UniformNadirScanning(value_per_second=1 / duration),\n", + " rewarder=data.ScanningTimeReward(),\n", + " time_limit=duration,\n", + " failure_penalty=-1.0,\n", + " terminate_on_time_limit=True,\n", + " log_level=\"WARNING\",\n", + ")\n", + "\n", + "obs_dim = env.observation_space.shape[0]\n", + "n_actions = env.action_space.n\n", + "print(\"Observation size:\", obs_dim)\n", + "print(\"Number of actions:\", n_actions)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The step duration is reported in the info dictionary on every step:" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "observation, info = env.reset(seed=SEED)\n", + "for action in range(n_actions):\n", + " observation, reward, terminated, truncated, info = env.step(action)\n", + " print(f\"Action {action}: d_ts = {info['d_ts']:.0f} s, reward = {reward:.4f}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Q-Network and Replay Buffer\n", + "The Q-network is a small multilayer perceptron that maps an observation to one Q-value per\n", + "action. The replay buffer stores the usual $(s, a, r, s', \\text{terminated})$ tuple **plus\n", + "the step duration** $\\Delta t$, which the time-discounted target needs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "class QNetwork(nn.Module):\n", + " def __init__(self, obs_dim: int, n_actions: int, hidden=(128, 128)):\n", + " super().__init__()\n", + " layers, in_dim = [], obs_dim\n", + " for width in hidden:\n", + " layers += [nn.Linear(in_dim, width), nn.ReLU()]\n", + " in_dim = width\n", + " layers.append(nn.Linear(in_dim, n_actions))\n", + " self.net = nn.Sequential(*layers)\n", + "\n", + " def forward(self, x: torch.Tensor) -> torch.Tensor:\n", + " return self.net(x)\n", + "\n", + "\n", + "class ReplayBuffer:\n", + " def __init__(self, capacity: int, obs_dim: int):\n", + " self.capacity = capacity\n", + " self.index = 0\n", + " self.size = 0\n", + " self.obs = np.zeros((capacity, obs_dim), dtype=np.float32)\n", + " self.next_obs = np.zeros((capacity, obs_dim), dtype=np.float32)\n", + " self.actions = np.zeros(capacity, dtype=np.int64)\n", + " self.rewards = np.zeros(capacity, dtype=np.float32)\n", + " self.terminated = np.zeros(capacity, dtype=np.float32)\n", + " self.d_ts = np.zeros(capacity, dtype=np.float32)\n", + "\n", + " def add(self, obs, action, reward, next_obs, terminated, d_t):\n", + " i = self.index\n", + " self.obs[i] = obs\n", + " self.actions[i] = action\n", + " self.rewards[i] = reward\n", + " self.next_obs[i] = next_obs\n", + " self.terminated[i] = terminated\n", + " self.d_ts[i] = d_t\n", + " self.index = (i + 1) % self.capacity\n", + " self.size = min(self.size + 1, self.capacity)\n", + "\n", + " def sample(self, batch_size: int):\n", + " idx = np.random.randint(0, self.size, size=batch_size)\n", + " return tuple(\n", + " torch.as_tensor(x[idx])\n", + " for x in (\n", + " self.obs,\n", + " self.actions,\n", + " self.rewards,\n", + " self.next_obs,\n", + " self.terminated,\n", + " self.d_ts,\n", + " )\n", + " )" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## DQN Agent\n", + "The agent keeps an online network and a slowly updated target network\n", + "([Polyak averaging](https://arxiv.org/abs/1509.02971) with rate ``tau``). It uses the Huber\n", + "loss and gradient clipping for stability. The only change from textbook DQN is the\n", + "discount: ``gamma ** d_t`` replaces a constant ``gamma``.\n", + "\n", + "With ``gamma=0.9995`` per second, a reward one orbit (about 5700 s) away is weighted by\n", + "roughly 0.06. That is enough lookahead for the agent to learn to charge and desaturate\n", + "*before* a failure occurs." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "class DQNAgent:\n", + " def __init__(\n", + " self,\n", + " obs_dim: int,\n", + " n_actions: int,\n", + " gamma: float = 0.9995, # per second\n", + " lr: float = 2.5e-4,\n", + " tau: float = 0.005,\n", + " batch_size: int = 64,\n", + " buffer_size: int = 50_000,\n", + " ):\n", + " self.n_actions = n_actions\n", + " self.gamma = gamma\n", + " self.tau = tau\n", + " self.batch_size = batch_size\n", + " self.q = QNetwork(obs_dim, n_actions)\n", + " self.q_target = QNetwork(obs_dim, n_actions)\n", + " self.q_target.load_state_dict(self.q.state_dict())\n", + " self.optimizer = torch.optim.Adam(self.q.parameters(), lr=lr)\n", + " self.buffer = ReplayBuffer(buffer_size, obs_dim)\n", + "\n", + " def act(self, obs, epsilon: float = 0.0) -> int:\n", + " if random.random() < epsilon:\n", + " return random.randrange(self.n_actions)\n", + " with torch.no_grad():\n", + " q_values = self.q(torch.as_tensor(obs, dtype=torch.float32))\n", + " return int(q_values.argmax())\n", + "\n", + " def update(self) -> float:\n", + " obs, actions, rewards, next_obs, terminated, d_ts = self.buffer.sample(\n", + " self.batch_size\n", + " )\n", + " with torch.no_grad():\n", + " next_q = self.q_target(next_obs).max(dim=1).values\n", + " # Semi-MDP target: discount by the elapsed time, not the step count\n", + " target = rewards + self.gamma**d_ts * (1.0 - terminated) * next_q\n", + " q = self.q(obs).gather(1, actions.unsqueeze(1)).squeeze(1)\n", + " loss = F.smooth_l1_loss(q, target)\n", + "\n", + " self.optimizer.zero_grad()\n", + " loss.backward()\n", + " nn.utils.clip_grad_norm_(self.q.parameters(), 10.0)\n", + " self.optimizer.step()\n", + "\n", + " with torch.no_grad():\n", + " for p_target, p in zip(self.q_target.parameters(), self.q.parameters()):\n", + " p_target.mul_(1.0 - self.tau).add_(self.tau * p)\n", + " return loss.item()\n", + "\n", + "\n", + "agent = DQNAgent(obs_dim, n_actions)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Training\n", + "The training loop is a standard epsilon-greedy loop through the Gymnasium API. Epsilon\n", + "decays linearly over the first half of training, and updates start once the buffer holds\n", + "``learning_starts`` transitions.\n", + "\n", + "Simulating a step costs much more than a gradient update on a small network, so the agent\n", + "takes ``updates_per_step = 4`` gradient steps for every environment step. Getting more\n", + "learning out of each simulated transition is the main practical advantage of an off-policy\n", + "method here.\n", + "\n", + "To keep the documentation build fast, the agent trains for only 20000 environment steps\n", + "(about 100 episodes, around ten minutes on a two-core machine). That is enough to see the agent learn to\n", + "avoid most failures, but far from convergence. For real results, train much longer and tune\n", + "the hyperparameters." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "total_steps = 20000\n", + "learning_starts = 1_000\n", + "updates_per_step = 4\n", + "epsilon_start, epsilon_end = 1.0, 0.05\n", + "epsilon_decay_steps = total_steps // 2\n", + "\n", + "episode_returns, episode_failures = [], []\n", + "step, episode = 0, 0\n", + "\n", + "while step < total_steps:\n", + " observation, info = env.reset(seed=SEED + episode)\n", + " episode += 1\n", + " done, episode_return = False, 0.0\n", + "\n", + " while not done:\n", + " fraction = min(1.0, step / epsilon_decay_steps)\n", + " epsilon = epsilon_start + fraction * (epsilon_end - epsilon_start)\n", + " action = agent.act(observation, epsilon)\n", + "\n", + " next_observation, reward, terminated, truncated, info = env.step(action)\n", + " agent.buffer.add(\n", + " observation, action, reward, next_observation, terminated, info[\"d_ts\"]\n", + " )\n", + " if agent.buffer.size >= learning_starts:\n", + " for _ in range(updates_per_step):\n", + " agent.update()\n", + "\n", + " observation = next_observation\n", + " episode_return += reward\n", + " step += 1\n", + " done = terminated or truncated\n", + "\n", + " episode_returns.append(episode_return)\n", + " episode_failures.append(not env.unwrapped.satellite.is_alive())\n", + " if episode % 10 == 0:\n", + " print(\n", + " f\"Episode {episode:3d} | step {step:6d} | epsilon {epsilon:.2f} | \"\n", + " f\"mean return (last 10) {np.mean(episode_returns[-10:]):+.3f} | \"\n", + " f\"failure rate (last 10) {np.mean(episode_failures[-10:]):.0%}\"\n", + " )" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The learning curve shows returns rising as the agent stops running out of battery or\n", + "saturating its reaction wheels. Failed episodes end early and incur the ``-1`` penalty." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "window = 10\n", + "moving_average = np.convolve(episode_returns, np.ones(window) / window, mode=\"valid\")\n", + "\n", + "fig, ax = plt.subplots(figsize=(7, 3.5))\n", + "ax.plot(episode_returns, alpha=0.3, label=\"Episode return\")\n", + "ax.plot(\n", + " np.arange(window - 1, len(episode_returns)), moving_average, label=\"10-episode mean\"\n", + ")\n", + "ax.axhline(0.0, color=\"k\", linewidth=0.5)\n", + "ax.set_xlabel(\"Episode\")\n", + "ax.set_ylabel(\"Return\")\n", + "ax.set_title(\"DQN training on the scanning/downlink environment\")\n", + "ax.legend()\n", + "plt.show()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluating the Policy\n", + "The greedy policy (no exploration) is compared with a random policy on episodes seeded\n", + "differently from training." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "def evaluate(policy, n_episodes=10, seed=1000):\n", + " returns, failures = [], []\n", + " for k in range(n_episodes):\n", + " observation, _ = env.reset(seed=seed + k)\n", + " done, episode_return = False, 0.0\n", + " while not done:\n", + " observation, reward, terminated, truncated, _ = env.step(\n", + " policy(observation)\n", + " )\n", + " episode_return += reward\n", + " done = terminated or truncated\n", + " returns.append(episode_return)\n", + " failures.append(not env.unwrapped.satellite.is_alive())\n", + " return np.mean(returns), np.std(returns), np.mean(failures)\n", + "\n", + "\n", + "for name, policy in [\n", + " (\"DQN (greedy)\", lambda o: agent.act(o)),\n", + " (\"Random\", lambda o: env.action_space.sample()),\n", + "]:\n", + " mean, std, failure_rate = evaluate(policy)\n", + " print(\n", + " f\"{name:13s} return {mean:+.3f} +/- {std:.3f} | failure rate {failure_rate:.0%}\"\n", + " )" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Next Steps\n", + "* **Save the policy** with ``torch.save(agent.q.state_dict(), \"dqn.pt\")``. At deployment\n", + " time, the greedy action is ``agent.q(obs).argmax()``.\n", + "* **Train longer**, and consider standard DQN extensions such as\n", + " [Double DQN](https://arxiv.org/abs/1509.06461), prioritized replay, or n-step returns.\n", + " For n-step returns, discount by the *sum* of the step durations.\n", + "* **Scale up throughput.** Steps are cheap to store but slow to simulate. For large runs,\n", + " collect experience from several environment copies in parallel, as the\n", + " [RLlib example](rllib_training.ipynb) does." + ] + } + ], + "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.10.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} From d8d459a640396bd5f3e6f505bba051bf88d03c49 Mon Sep 17 00:00:00 2001 From: Vishwa GW <136871744+vishwagw@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:32:26 +0530 Subject: [PATCH 2/3] Add DQN example to examples index --- examples/_default.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/_default.rst b/examples/_default.rst index a8945013..71b64c63 100644 --- a/examples/_default.rst +++ b/examples/_default.rst @@ -36,7 +36,8 @@ Training .. toctree:: :maxdepth: 1 - rllib_training + rllib_training + dqn_training time_discounted_gae async_multiagent_training curriculum_learning From 0ad238a5b3b9cc1c8ec3cfb9416d1fb7bba37576 Mon Sep 17 00:00:00 2001 From: Vishwa GW <136871744+vishwagw@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:35:51 +0530 Subject: [PATCH 3/3] Add release note for DQN example --- docs/source/release_notes.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst index 7d62bbde..a713efcd 100644 --- a/docs/source/release_notes.rst +++ b/docs/source/release_notes.rst @@ -14,6 +14,8 @@ Development - |version| overlapping assignments and passive target spacecraft. Communication shares channel data while preserving local data production and reward attribution. See :doc:`the mixed ground and space imaging example `. +* Add :doc:`the DQN training example `, which trains a DQN agent + in plain PyTorch using semi-MDP time discounting (``gamma ** d_ts``) in the TD target. * Add :doc:`the space-to-space imaging example ` for independently orbiting spacecraft: replayable RSO catalogs, explicitly named observers, moving-target access,