diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst index eff3ea94..fd09b8fc 100644 --- a/docs/source/release_notes.rst +++ b/docs/source/release_notes.rst @@ -10,6 +10,8 @@ Development - |version| corrupted ``bisect``-based opportunity lookups. See issue #205. * Fix a bug where taking the charge action would not work properly because of a missing message subscription in the :class:`~bsk_rl.sim.fsw.BasicFSWModel`. +* Update `Cloud Environment with Re-imaging `_ example + adding a new network architecture and example heuristics. Version 1.3.0 ------------- diff --git a/examples/cloud_environment_with_reimaging.ipynb b/examples/cloud_environment_with_reimaging.ipynb index be523e0b..a6d15a38 100644 --- a/examples/cloud_environment_with_reimaging.ipynb +++ b/examples/cloud_environment_with_reimaging.ipynb @@ -5,7 +5,7 @@ "metadata": {}, "source": [ "# Cloud Environment with Re-imaging\n", - "This tutorial demonstrates the configuration and use of a BSK-RL environment considering cloud coverage and re-imaging capabilities. Two reward functions are presented: a single-picture binary case (where targets are deemed occluded by clouds or not and no re-imaging is allowed) and a re-imaging case where the problem is formulated in terms of the targets' probability of being successfully observed. Still, the satellite cannot observe the true cloud coverage of each target, only its forecast. The satellite has to image targets while keeping a positive battery level. This example script is part of an upcoming publication.\n", + "This tutorial demonstrates the configuration and use of a BSK-RL environment considering cloud coverage and re-imaging capabilities. Two reward functions are introduced: a single-picture binary case (where targets are deemed occluded by clouds or not and no re-imaging is allowed) and a re-imaging case where the problem is formulated in terms of the targets' probability of being successfully observed. Still, the satellite cannot observe the true cloud coverage of each target, only its forecast. The satellite has to image targets while keeping a positive battery level. Moreover, custom actor and critic modules are defined leveraging a token-based architecture. This example script is part of an upcoming publication.\n", "\n", "## Loading Modules" ] @@ -23,7 +23,7 @@ "from Basilisk.architecture import bskLogging\n", "from Basilisk.utilities import orbitalMotion\n", "from bsk_rl import act, obs, sats\n", - "from bsk_rl.sim import dyn, fsw\n", + "from bsk_rl.sim import fsw\n", "from bsk_rl.scene.targets import UniformTargets\n", "from bsk_rl.data.base import Data, DataStore, GlobalReward\n", "from bsk_rl.data.unique_image_data import (\n", @@ -111,7 +111,7 @@ " b_S0 = 1 - b_S1\n", " target.belief = np.array([b_S0, b_S1])\n", " target.prev_obs = -np.random.uniform(\n", - " self.prev_obs_init[0], self.prev_obs_init[0]\n", + " self.prev_obs_init[0], self.prev_obs_init[1]\n", " )\n", " target.belief_update_var = 0.0\n", "\n", @@ -133,12 +133,12 @@ "\n", "$$\n", "R = \\begin{cases}\n", - "\\rho_i & \\text{if } c_{p_i} \\leq c_{\\text{thr}_i} \\\\\n", + "r_i & \\text{if } c_{t_i} \\leq c_{\\text{thr}_i} \\\\\n", "0 & \\text{otherwise.}\n", "\\end{cases}\n", "$$\n", "\n", - "where $\\rho_i$ is priority, $c_{p_i}$ is the true cloud coverage, and $c_{\\text{thr}_i}$ is the `reward_threshold` for target $i$. For a case where the reward is linearly proportional to the cloud coverage, see [Cloud Environment](../examples/cloud_environment.rst)" + "where $r_i$ is priority, $c_{t_i}$ is the true cloud coverage, and $c_{\\text{thr}_i}$ is the `reward_threshold` for target $i$. For a case where the reward is linearly proportional to the cloud coverage, see [Cloud Environment](../examples/cloud_environment.rst)" ] }, { @@ -312,8 +312,8 @@ "\n", "$$\n", "R = \\begin{cases}\n", - "\\rho_i\\alpha_i\\Delta \\text{P}(S=1) + \\rho_i(1 - \\alpha) & \\text{ if } \\text{P}_i(S=1) \\geq \\theta_{\\text{thr}_i} \\\\\n", - "\\rho_i\\alpha_i\\Delta \\text{P}(S=1) & \\text{ otherwise.}\n", + "r_i\\alpha_i\\Delta \\text{P}(S=1) + r_i(1 - \\alpha) & \\text{ if } \\text{P}_i(S=1) \\geq \\theta_{\\text{thr}_i} \\\\\n", + "r_i\\alpha_i\\Delta \\text{P}(S=1) & \\text{ otherwise.}\n", "\\end{cases}\n", "$$\n" ] @@ -514,13 +514,13 @@ " for target, belief_variation in zip(\n", " new_data.imaged, new_data.list_belief_update_var\n", " ):\n", - " reward[sat_id] += self.reward_fn(\n", - " target.priority, belief_variation, self.alpha, reach_threshold=False\n", - " )\n", - " for target in new_data.imaged_complete:\n", - " reward[sat_id] += self.reward_fn(\n", - " target.priority, None, self.alpha, reach_threshold=True\n", - " )\n", + " if target not in self.data.imaged_complete:\n", + " reward[sat_id] += self.reward_fn(\n", + " target.priority,\n", + " belief_variation,\n", + " self.alpha,\n", + " reach_threshold=target in new_data.imaged_complete,\n", + " )\n", " return reward\n", "\n", "\n", @@ -561,7 +561,7 @@ "The update in the success probability is given by:\n", "\n", "$$\n", - "\\text{P}^{(k+1)}(S=1) = 1 - \\text{P}^{(k)}(S=1)\\bar{c}_{f_i}\n", + "\\text{P}^{(k+1)}(S=1) = 1 - \\text{P}^{(k)}(S=0)\\bar{c}_{f_i}\n", "$$\n", "\n", "To penalize two consecutive pictures without enough elapsed time (and not enough shift in clouds' position), a new cloud-free probability variable $g_{f_i}$ is introduced such that\n", @@ -668,14 +668,13 @@ "The satellite has observations and actions associated with it that are relevant to the decision-making process. The observation space can be modified to include information about the targets and the weather (cloud coverage forecast, reward threshold, success probability, etc) which allows better informed decision-making.\n", "\n", "* [Observations](../api_reference/obs/index.rst): \n", - " - SatProperties: Body angular velocity, instrument pointing direction, body position, body velocity, battery charge (properties in [flight software model](../api_reference/sim/fsw/index.rst) or [dynamics model](../api_reference/sim/dyn/index.rst)). Also, customized dynamics property in CustomDynModel below: Angle between the sun and the solar panel.\n", - " - OpportunityProperties: Target's priority, cloud coverage forecast, standard deviation of cloud coverage forecast, probability of being successfully imaged, and last time it was imaged (upcoming 32 targets). \n", - " - Time: Simulation time.\n", + " - SatProperties: Body angular velocity, instrument pointing direction, body position, body velocity, battery charge (properties in [flight software model](../api_reference/sim/fsw/index.rst) or [dynamics model](../api_reference/sim/dyn/index.rst)). Also, customized dynamics property in CustomDynModel below: spacecraft-to-Sun unit direction vector.\n", + " - OpportunityProperties: Target's priority, cloud coverage forecast, standard deviation of cloud coverage forecast, probability of being successfully imaged, and last time it was imaged (upcoming 40 targets). \n", " - Eclipse: Next eclipse start and end times. \n", "* [Actions](../api_reference/act/index.rst):\n", " - Charge: Enter a sun-pointing charging mode for 60 seconds.\n", - " - Image: Image target from upcoming 32 targets\n", - "* [Dynamics model](../api_reference/sim/dyn/index.rst): FullFeaturedDynModel is used and a property, angle between sun and solar panel, is added.\n", + " - Image: Image target from upcoming 40 targets\n", + "* [Dynamics model](../api_reference/sim/dyn/index.rst): FullFeaturedDynModel is used and a property, spacecraft-to-Sun unit direction vector, is added.\n", "* [Flight software model](../api_reference/sim/fsw/index.rst): SteeringImagerFSWModel is used." ] }, @@ -685,68 +684,154 @@ "metadata": {}, "outputs": [], "source": [ - "class CustomSatComposed(sats.ImagingSatellite):\n", + "def s_hat_H(sat):\n", + " r_SN_N = (\n", + " sat.simulator.world.gravFactory.spiceObject.planetStateOutMsgs[\n", + " sat.simulator.world.sun_index\n", + " ]\n", + " .read()\n", + " .PositionVector\n", + " )\n", + " r_BN_N = sat.dynamics.r_BN_N\n", + " r_SB_N = np.array(r_SN_N) - np.array(r_BN_N)\n", + " r_SB_H = sat.dynamics.HN @ r_SB_N\n", + " return r_SB_H / np.linalg.norm(r_SB_H)\n", + "\n", + "\n", + "class CloudCoverDensity(obs.Observation):\n", + " def __init__(\n", + " self,\n", + " interval_duration=60 * 3,\n", + " intervals=10,\n", + " norm=3,\n", + " ):\n", + " self.satellite: \"sats.ImagingSatellite\"\n", + " super().__init__()\n", + " self.interval_duration = interval_duration\n", + " self.intervals = intervals\n", + " self.norm = norm\n", + "\n", + " def get_obs(self):\n", + " if self.intervals == 0:\n", + " return []\n", + "\n", + " self.satellite.calculate_additional_windows(\n", + " self.simulator.sim_time\n", + " + (self.intervals + 1) * self.interval_duration\n", + " - self.satellite.window_calculation_time\n", + " )\n", + " soonest = self.satellite.upcoming_opportunities_dict(types=\"target\")\n", + " cloud_covers = np.array([target.cloud_cover_forecast for target in soonest])\n", + " times = np.array([opportunities[0][1] for opportunities in soonest.values()])\n", + " time_bins = np.floor((times - self.simulator.sim_time) / self.interval_duration)\n", + " densities = [sum(cloud_covers[time_bins == i]) for i in range(self.intervals)]\n", + " return np.array(densities) / self.norm\n", + "\n", + "\n", + "class CS_F(sats.ImagingSatellite):\n", " observation_spec = [\n", " obs.SatProperties(\n", - " dict(prop=\"omega_BP_P\", norm=0.03),\n", - " dict(prop=\"c_hat_P\"),\n", - " dict(prop=\"r_BN_P\", norm=orbitalMotion.REQ_EARTH * 1e3),\n", - " dict(prop=\"v_BN_P\", norm=7616.5),\n", - " dict(prop=\"battery_charge_fraction\"),\n", - " dict(prop=\"solar_angle_norm\"),\n", + " dict(prop=\"omega_BN_B\", norm=0.03), # 0,1,2 - CHECKED\n", + " dict(prop=\"c_hat_H\"), # 3,4,5 - CHECKED\n", + " dict(prop=\"r_BN_P\", norm=orbitalMotion.REQ_EARTH * 1e3), # 6,7,8 - CHECKED\n", + " dict(prop=\"v_BN_P\", norm=7616.5), # 9,10,11 - CHECKED\n", + " dict(prop=\"battery_charge_fraction\"), # 12 - CHECKED\n", + " dict(prop=\"s_hat_H\", fn=s_hat_H), # 17, 18, 19 - CHECKED\n", " ),\n", - " obs.Eclipse(),\n", + " obs.Eclipse(norm=5700), # 22,23 - CHECKED\n", + " CloudCoverDensity(intervals=20, norm=5),\n", " obs.OpportunityProperties(\n", - " dict(prop=\"priority\"),\n", + " dict(prop=\"priority\"), # - CHECKED\n", + " dict(prop=\"r_LB_H\", norm=orbitalMotion.REQ_EARTH * 1e3), # - CHECKED\n", + " dict(\n", + " prop=\"target_angle\", norm=np.pi / 2\n", + " ), # - CHECKED (LOOK IF THIS IS 1 OR 3)\n", + " dict(\n", + " prop=\"target_angle_rate\", norm=0.03\n", + " ), # - CHECKED (LOOK IF THIS IS 1 OR 3)\n", + " dict(prop=\"opportunity_open\", norm=300.0), # - CHECKED\n", + " dict(prop=\"opportunity_close\", norm=300.0), # - CHECKED\n", " dict(\n", - " fn=lambda sat, opp: opp[\"object\"].cloud_cover_forecast\n", + " prop=\"cloud_forecast\",\n", + " fn=lambda sat, opp: opp[\"object\"].cloud_cover_forecast,\n", " ), # Cloud coverage forecast (percentage of the area covered by clouds)\n", " dict(\n", - " fn=lambda sat, opp: opp[\"object\"].cloud_cover_sigma\n", + " prop=\"cloud_sigma\",\n", + " fn=lambda sat, opp: opp[\"object\"].cloud_cover_sigma,\n", + " norm=0.05,\n", " ), # Confidence on the cloud coverage forecast\n", - " # dict(fn=lambda sat, opp: opp[\"object\"].reward_threshold), #Reward threshold for each target. Uncomment if using variable threshold\n", - " dict(\n", - " fn=lambda sat, opp: opp[\"object\"].belief[1]\n", - " ), # Probability of successfully imaging the target. Used only in the re-imaging case\n", + " type=\"target\",\n", + " n_ahead_observe=40,\n", + " ),\n", + " ]\n", + "\n", + " action_spec = [\n", + " act.Charge(duration=60.0),\n", + " act.Image(n_ahead_image=40),\n", + " ]\n", + "\n", + " fsw_type = fsw.SteeringImagerFSWModel\n", + "\n", + "\n", + "def time_since_prev_obs(sat, opp):\n", + " prev_obs = opp[\"object\"].prev_obs\n", + " if prev_obs is None:\n", + " return 0.0\n", + " else:\n", + " return sat.simulator.sim_time - prev_obs\n", + "\n", + "\n", + "def belief_expected(sat, opp):\n", + " belief = opp[\"object\"].belief.copy()\n", + " cloud_cover_forecast = opp[\"object\"].cloud_cover_forecast.copy()\n", + " delta_t = time_since_prev_obs(sat, opp)\n", + " expected = belief_update_func(belief, cloud_cover_forecast, delta_t)\n", + "\n", + " return expected[1]\n", + "\n", + "\n", + "class CS_F_Bayesian(sats.ImagingSatellite):\n", + " observation_spec = [\n", + " obs.SatProperties(\n", + " dict(prop=\"omega_BN_B\", norm=0.03),\n", + " dict(prop=\"c_hat_H\"),\n", + " dict(prop=\"r_BN_P\", norm=orbitalMotion.REQ_EARTH * 1e3),\n", + " dict(prop=\"v_BN_P\", norm=7616.5),\n", + " dict(prop=\"battery_charge_fraction\"),\n", + " dict(prop=\"s_hat_H\", fn=s_hat_H),\n", + " ),\n", + " obs.Eclipse(norm=5700),\n", + " CloudCoverDensity(intervals=20, norm=5),\n", + " obs.OpportunityProperties(\n", + " dict(prop=\"priority\"),\n", + " dict(prop=\"r_LB_H\", norm=orbitalMotion.REQ_EARTH * 1e3),\n", + " dict(prop=\"target_angle\", norm=np.pi / 2),\n", + " dict(prop=\"target_angle_rate\", norm=0.03),\n", + " dict(prop=\"opportunity_open\", norm=300.0),\n", + " dict(prop=\"opportunity_close\", norm=300.0),\n", " dict(\n", - " fn=lambda sat, opp: opp[\"object\"].prev_obs, norm=5700\n", - " ), # Previous observation time. Used only in the re-imaging case\n", + " prop=\"cloud_forecast\",\n", + " fn=lambda sat, opp: opp[\"object\"].cloud_cover_forecast,\n", + " ),\n", + " dict(prop=\"belief\", fn=lambda sat, opp: opp[\"object\"].belief[1], norm=1.0),\n", + " dict(prop=\"time_since_prev_obs\", fn=time_since_prev_obs, norm=5700.0),\n", + " dict(prop=\"belief_expected\", fn=belief_expected, norm=1.0),\n", " type=\"target\",\n", - " n_ahead_observe=32,\n", + " n_ahead_observe=40,\n", " ),\n", - " obs.Time(),\n", " ]\n", "\n", " action_spec = [\n", " act.Charge(duration=60.0),\n", - " act.Image(n_ahead_image=32),\n", + " act.Image(n_ahead_image=40),\n", " ]\n", "\n", - " # Modified the constructor to include the belief update function\n", + " fsw_type = fsw.SteeringImagerFSWModel\n", + "\n", " def __init__(self, *args, belief_update_func=None, **kwargs) -> None:\n", " super().__init__(*args, **kwargs)\n", " self.belief_update_func = belief_update_func\n", - "\n", - " class CustomDynModel(dyn.FullFeaturedDynModel):\n", - " @property\n", - " def solar_angle_norm(self) -> float:\n", - " sun_vec_N = (\n", - " self.world.gravFactory.spiceObject.planetStateOutMsgs[\n", - " self.world.sun_index\n", - " ]\n", - " .read()\n", - " .PositionVector\n", - " )\n", - " sun_vec_N_hat = sun_vec_N / np.linalg.norm(sun_vec_N)\n", - " solar_panel_vec_B = np.array([0, 0, -1]) # Not default configuration\n", - " mat = np.transpose(self.BN)\n", - " solar_panel_vec_N = np.matmul(mat, solar_panel_vec_B)\n", - " error_angle = np.arccos(np.dot(solar_panel_vec_N, sun_vec_N_hat))\n", - "\n", - " return error_angle / np.pi\n", - "\n", - " dyn_type = CustomDynModel\n", - " fsw_type = fsw.SteeringImagerFSWModel" + " self.also_image_filter = []" ] }, { @@ -786,14 +871,16 @@ "metadata": {}, "outputs": [], "source": [ - "dataStorageCapacity = 20 * 8e6 * 100\n", - "sat_args = CustomSatComposed.default_sat_args(\n", + "dataStorageCapacity: float = 20 * 8e6 * 100\n", + "sat_args = CS_F_Bayesian.default_sat_args(\n", " imageAttErrorRequirement=0.01,\n", " imageRateErrorRequirement=0.01,\n", " batteryStorageCapacity=80.0 * 3600 * 2,\n", " storedCharge_Init=lambda: np.random.uniform(0.4, 1.0) * 80.0 * 3600 * 2,\n", - " u_max=0.2,\n", - " K1=0.5,\n", + " u_max=0.1,\n", + " K1=0.2,\n", + " K3=0.5,\n", + " servo_P=30,\n", " nHat_B=np.array([0, 0, -1]),\n", " imageTargetMinimumElevation=np.radians(45),\n", " rwBasePower=20,\n", @@ -806,6 +893,11 @@ " -1, 1, 3\n", " ), # Initialize reaction wheel speeds close to zero\n", " dataStorageCapacity=dataStorageCapacity, # Large storage to avoid filling up in three orbits\n", + " # oe=partial(\n", + " # random_circular_orbit,\n", + " # i=45.0,\n", + " # alt=500,\n", + " # ), # Optional for when using a single satellite\n", ")" ] }, @@ -834,7 +926,7 @@ ")\n", "\n", "satellites = [\n", - " CustomSatComposed(f\"EO-{i}\", sat_args, belief_update_func=belief_update_func)\n", + " CS_F_Bayesian(f\"EO-{i}\", sat_args, belief_update_func=belief_update_func)\n", " for i in range(5)\n", "]\n", "\n", @@ -850,7 +942,7 @@ " sim_rate=0.5,\n", " max_step_duration=300.0,\n", " time_limit=95 * 60 / 2, # half orbit\n", - " log_level=\"INFO\",\n", + " log_level=\"WARNING\",\n", " failure_penalty=0.0,\n", " # disable_env_checker=True, # For debugging\n", ")" @@ -888,10 +980,217 @@ "print(\"Actions:\", env.satellites[0].action_description, \"\\n\")\n", "print(\"States:\", env.unwrapped.satellites[0].observation_description, \"\\n\")\n", "\n", - "# Using the composed satellite features also provides a human-readable state:\n", - "for satellite in env.unwrapped.satellites:\n", - " for k, v in satellite.observation_builder.obs_dict().items():\n", - " print(f\"{k}: {v}\")" + "# # Uncomment to see the observation for each satellite\n", + "# # Using the composed satellite features also provides a human-readable state:\n", + "# for satellite in env.unwrapped.satellites:\n", + "# for k, v in satellite.observation_builder.obs_dict().items():\n", + "# print(f\"{k}: {v}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Greedy Heuristic\n", + "\n", + "A greedy heuristic can be created to task the satellite. For the heuristic with $\\alpha=1.0$, a target is selected from the $N=35$ upcoming targets with\n", + "$$\n", + "\\rho^\\star=\\arg\\max_{\\rho_i \\in \\mathcal{U}_{35}} \\left\\{\\frac{r_i \\Delta\\text{P}_i(S=1)}{t_{i,{\\text{max}}}} \\; : \\; \\text{P}_i(S=1) < \\theta_{\\text{thr}_i} \\; t_{i,\\text{max}} < t_{i,\\text{close}}\\right\\}\n", + "$$\n", + "assuming that the next image is taken at the next decision step. For the heuristic with $\\alpha=0.0$ the target is selected with\n", + "$$\n", + " \\rho^\\star=\\arg\\max_{\\rho_i \\in \\mathcal{U}_{35}} \\left\\{\\frac{r_i}{t_{i,{\\text{max}}}k_i}, \\; : \\; \\text{P}_i(S=1) < \\theta_{\\text{thr}_i}, \\; t_{i,\\text{max}} < t_{i,\\text{close}}\\right\\}\n", + "$$\n", + "where $k_i$ estimates the number of images required to reach the success threshold and assuming enough elapsed time between images." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def slew_time_regression(\n", + " tgt_angle: float,\n", + " tgt_angle_rate: float,\n", + " x_0: float = 3.504,\n", + " x_1: float = 96.8639,\n", + " x_2: float = 1325.3045,\n", + " x_3: float = -23.6751,\n", + " x_4: float = -1.681e4,\n", + " x_5: float = -1009.7423,\n", + ") -> float:\n", + " \"\"\"\n", + " Regression model to estimate the slew time based on the target angle and target angle rate.\n", + "\n", + " Created by fitting a polynomial regression model to the data obtained from the simulation. A NN could be used instead for better accuracy.\n", + " \"\"\"\n", + " slew_time = (\n", + " x_0\n", + " + x_1 * tgt_angle\n", + " + x_2 * tgt_angle_rate\n", + " + x_3 * tgt_angle**2\n", + " + x_4 * tgt_angle_rate**2\n", + " + x_5 * tgt_angle * tgt_angle_rate\n", + " )\n", + " return np.clip(slew_time, 1e-3, None)\n", + "\n", + "\n", + "def collect_tgt_info(sat, lookahead: int) -> tuple:\n", + " \"\"\"\n", + " Collects information about the targets in the satellite's observation space.\n", + " \"\"\"\n", + "\n", + " obs = sat.get_obs()\n", + "\n", + " list_slew_time = []\n", + " list_priority = []\n", + " list_cloud_forecast = []\n", + " list_window_open = []\n", + " list_window_close = []\n", + " list_cloud_sigma = []\n", + " list_belief = []\n", + " list_time_since_prev_obs = []\n", + "\n", + " obs_keys = sat.observation_builder.obs_array_keys()\n", + "\n", + " for tgt_i in range(lookahead):\n", + " pos_angle_i = obs_keys.index(f\"target.target_{tgt_i}.target_angle_normd\")\n", + " pos_angle_rate_i = obs_keys.index(\n", + " f\"target.target_{tgt_i}.target_angle_rate_normd\"\n", + " )\n", + " pos_window_open_i = obs_keys.index(\n", + " f\"target.target_{tgt_i}.opportunity_open_normd\"\n", + " )\n", + " pos_window_close_i = obs_keys.index(\n", + " f\"target.target_{tgt_i}.opportunity_close_normd\"\n", + " )\n", + " pos_priority_i = obs_keys.index(f\"target.target_{tgt_i}.priority\")\n", + " pos_cloud_forecast_i = obs_keys.index(f\"target.target_{tgt_i}.cloud_forecast\")\n", + " if f\"target.target_{tgt_i}.cloud_sigma_normd\" in obs_keys:\n", + " pos_cloud_sigma_i = obs_keys.index(\n", + " f\"target.target_{tgt_i}.cloud_sigma_normd\"\n", + " )\n", + " else:\n", + " pos_cloud_sigma_i = None\n", + " if f\"target.target_{tgt_i}.belief\" in obs_keys:\n", + " pos_belief_i = obs_keys.index(f\"target.target_{tgt_i}.belief\")\n", + " pos_time_since_prev_obs_i = obs_keys.index(\n", + " f\"target.target_{tgt_i}.time_since_prev_obs_normd\"\n", + " )\n", + " else:\n", + " pos_belief_i = None\n", + " pos_time_since_prev_obs_i = None\n", + "\n", + " angle_i = obs[pos_angle_i] * np.pi / 2\n", + " angle_rate_i = obs[pos_angle_rate_i] * 0.03\n", + " window_open_i = obs[pos_window_open_i] * 300.0\n", + " window_close_i = obs[pos_window_close_i] * 300.0\n", + " priority_i = obs[pos_priority_i]\n", + " cloud_forecast = obs[pos_cloud_forecast_i]\n", + " cloud_sigma_i = (\n", + " obs[pos_cloud_sigma_i] * 0.05 if pos_cloud_sigma_i is not None else 0\n", + " )\n", + " belief_i = obs[pos_belief_i] if pos_belief_i is not None else 0\n", + " time_since_prev_obs_i = (\n", + " obs[pos_time_since_prev_obs_i] * 5700.0\n", + " if pos_time_since_prev_obs_i is not None\n", + " else 0\n", + " )\n", + "\n", + " slew_time_i = slew_time_regression(angle_i, angle_rate_i)\n", + "\n", + " list_slew_time.append(slew_time_i)\n", + " list_priority.append(priority_i)\n", + " list_cloud_forecast.append(cloud_forecast)\n", + " list_window_open.append(window_open_i)\n", + " list_window_close.append(window_close_i)\n", + " list_cloud_sigma.append(cloud_sigma_i)\n", + " list_belief.append(belief_i)\n", + " list_time_since_prev_obs.append(time_since_prev_obs_i)\n", + "\n", + " return (\n", + " np.array(list_slew_time),\n", + " np.array(list_priority),\n", + " np.array(list_cloud_forecast),\n", + " np.array(list_window_open),\n", + " np.array(list_window_close),\n", + " np.array(list_cloud_sigma),\n", + " np.array(list_belief),\n", + " np.array(list_time_since_prev_obs),\n", + " )\n", + "\n", + "\n", + "def greedy_heuristic_reimaging(sat, alpha, lookahead):\n", + " \"\"\"\n", + " Greedy heuristic for reimaging targets based on the expected belief update and the priority of the target.\n", + "\n", + " Args:\n", + " sat: Satellite object.\n", + " alpha: Tuning parameter. Should be 0 or 1.\n", + " lookahead: Number of targets to consider in the observation space.\n", + "\n", + " Returns:\n", + " int: Index of the best target to image next in the action space (0 action is charge).\n", + " \"\"\"\n", + "\n", + " (\n", + " list_slew_time,\n", + " list_priority,\n", + " list_cloud_forecast,\n", + " list_opening_windows,\n", + " list_closing_windows,\n", + " _,\n", + " list_belief,\n", + " list_time_since_prev_obs,\n", + " ) = collect_tgt_info(sat, lookahead)\n", + "\n", + " time_to_tgt = np.maximum(list_slew_time, list_opening_windows)\n", + "\n", + " time_to_tgt = np.clip(time_to_tgt, 1.0, None)\n", + "\n", + " list_scores = []\n", + " for i in range(lookahead):\n", + " if time_to_tgt[i] > list_closing_windows[i]:\n", + " list_scores.append(-np.inf)\n", + " else:\n", + " tgt_cloud_forecast = list_cloud_forecast[i]\n", + " tgt_belief = list_belief[i]\n", + " priority = list_priority[i]\n", + "\n", + " expected_belief = belief_update_func(\n", + " [1 - tgt_belief, tgt_belief],\n", + " tgt_cloud_forecast,\n", + " list_time_since_prev_obs[i],\n", + " )\n", + "\n", + " if alpha == 1:\n", + " score = (expected_belief[1] - tgt_belief) * priority / time_to_tgt[i]\n", + " else:\n", + " if expected_belief[1] >= 0.95:\n", + " score = priority / time_to_tgt[i]\n", + " else:\n", + " if tgt_cloud_forecast <= 1.0:\n", + " new_belief_temp = expected_belief\n", + " count = 1\n", + " for _ in range(10):\n", + " new_belief_temp = belief_update_func(\n", + " new_belief_temp.copy(),\n", + " tgt_cloud_forecast,\n", + " list_time_since_prev_obs[i],\n", + " )\n", + " count += 1\n", + " if new_belief_temp[1] >= 0.95:\n", + " break\n", + " score = priority / time_to_tgt[i] / count\n", + " else:\n", + " score = 0.0\n", + "\n", + " list_scores.append(score)\n", + "\n", + " best_target = np.argmax(list_scores)\n", + "\n", + " return best_target + 1" ] }, { @@ -907,20 +1206,23 @@ "metadata": {}, "outputs": [], "source": [ - "count = 0\n", + "steps = 0\n", "while True:\n", - " if count == 0:\n", + " if steps == 0:\n", " # Vector with an action for each satellite (we can pass different actions for each satellite)\n", " # Tasking all satellites to charge (tasking None as the first action will raise a warning)\n", " action_dict = {sat_i.name: 0 for sat_i in env.satellites}\n", " else:\n", - " # Tasking random actions\n", - " action_dict = {sat_i.name: np.random.randint(0, 32) for sat_i in env.satellites}\n", - " count += 1\n", + " # Using the greedy heuristic to select the best target for each satellite based on the current observation\n", + " action_dict = {\n", + " sat_i.name: greedy_heuristic_reimaging(sat_i, 0.0, 35)\n", + " for sat_i in env.satellites\n", + " }\n", + " steps += 1\n", "\n", " observation, reward, terminated, truncated, info = env.step(action_dict)\n", "\n", - " if all(terminated.values()) or all(truncated.values()):\n", + " if all(terminated.values()) or all(truncated.values()) or steps >= 5:\n", " print(\"Episode complete.\")\n", " break" ] @@ -959,13 +1261,732 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Check [Training with RLlib PPO](../examples/rllib_training.ipynb) for an example on how to train the agent in this environment." + "A similar heuristic can be used for the single-picture case, where the expected reward is used instead of the rewards to compute the reward density per target with\n", + "\n", + "$$\n", + "\\rho_i^\\star = \\arg\\max_{\\rho_i \\in \\mathcal{U}_{35}} \\left\\{ \\frac{r_i p_i}{t_{i,\\text{max}}} \\; : \\; t_{i,\\text{max}} \\leq t_{i,\\mathrm{close}} \\right\\}\n", + "$$\n", + "\n", + "where the expected reward is computed with\n", + "\n", + "$$\n", + "p_i = \\text{P}(c_{t_i}\\leq c_{\\text{thr}_i}|c_{f_i}, \\sigma_{i})\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import scipy as sp\n", + "\n", + "\n", + "def greedy_heuristic_single_picture(sat, lookahead: int = 35):\n", + " \"\"\"\n", + " Greedy heuristic for selecting the best target to image based on the expected reward per time.\n", + "\n", + " Args:\n", + " sat: Satellite object.\n", + " lookahead: Number of targets to consider in the observation space.\n", + " Returns:\n", + " int: Index of the best target to image next in the action space (0 action is charge).\n", + " \"\"\"\n", + "\n", + " (\n", + " list_slew_time,\n", + " list_priority,\n", + " list_cloud_forecast,\n", + " list_opening_windows,\n", + " list_closing_windows,\n", + " list_cloud_sigma,\n", + " _,\n", + " _,\n", + " ) = collect_tgt_info(sat, lookahead)\n", + "\n", + " time_to_tgt = np.maximum(list_slew_time, list_opening_windows)\n", + "\n", + " time_to_tgt = np.clip(time_to_tgt, 1.0, None)\n", + "\n", + " list_scores = np.zeros(len(time_to_tgt))\n", + " for i in range(len(time_to_tgt)):\n", + " if time_to_tgt[i] > list_closing_windows[i]:\n", + " list_scores[i] = -np.inf\n", + " else:\n", + " tgt_probability_clear_i = np.clip(\n", + " sp.stats.norm.cdf(\n", + " 0.2,\n", + " loc=list_cloud_forecast[i],\n", + " scale=list_cloud_sigma[i],\n", + " ),\n", + " 0.0,\n", + " 1.0,\n", + " )\n", + " list_scores[i] = list_priority[i] * tgt_probability_clear_i / time_to_tgt[i]\n", + "\n", + " target_pos = np.argmax(list_scores)\n", + "\n", + " return target_pos + 1" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Training and custom actor-critic modules\n", + "\n", + "The observation space contains concatenated information regarding the spacecraft and individual targets, which naturally decomposes into two components: a spacecraft state vector and a set of target-specific state vectors. This representation motivates the use of shared target encoders, attention mechanisms, and permutation-invariant pooling operations instead of applying an MLP directly to the flattened observation. \n", + "\n", + "Then, custom actor and critic modules are defined for training with PPO." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Dict\n", + "\n", + "from ray.rllib.core import Columns\n", + "from ray.rllib.core.models.base import ACTOR, CRITIC, ENCODER_OUT\n", + "from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule\n", + "from ray.rllib.utils.annotations import (\n", + " override,\n", + ")\n", + "from ray.rllib.utils.framework import try_import_torch\n", + "from ray.rllib.utils.typing import TensorType\n", + "from ray.rllib.models.torch.torch_distributions import (\n", + " TorchCategorical,\n", + ")\n", + "from ray.rllib.core.models.configs import RecurrentEncoderConfig\n", + "from ray.rllib.algorithms.ppo.torch.ppo_torch_rl_module import PPOTorchRLModule\n", + "\n", + "torch, nn = try_import_torch()\n", + "\n", + "\n", + "class MultiHeadSelfAttention(nn.Module):\n", + " def __init__(self, d_in: int, d_model: int, n_heads: int = 4):\n", + " super().__init__()\n", + " assert d_model % n_heads == 0\n", + " self.n_heads = n_heads\n", + " self.d_head = d_model // n_heads\n", + "\n", + " self.qkv = nn.Linear(d_in, 3 * d_model, bias=False)\n", + " self.out = nn.Linear(d_model, d_in, bias=False)\n", + "\n", + " def forward(self, x):\n", + " B, N, _ = x.shape\n", + "\n", + " qkv = self.qkv(x)\n", + " qkv = qkv.view(B, N, 3, self.n_heads, self.d_head)\n", + " qkv = qkv.permute(2, 0, 3, 1, 4)\n", + "\n", + " q, k, v = qkv[0], qkv[1], qkv[2]\n", + "\n", + " attn = torch.nn.functional.scaled_dot_product_attention(\n", + " q, k, v, dropout_p=0.0, is_causal=False\n", + " )\n", + "\n", + " attn = attn.transpose(1, 2).contiguous().view(B, N, -1)\n", + " return self.out(attn)\n", + "\n", + "\n", + "class MultiHeadAttention(nn.Module):\n", + " def __init__(self, d_in_q: int, d_in_kv: int, d_model: int, n_heads: int = 4):\n", + " super().__init__()\n", + " assert d_model % n_heads == 0\n", + " self.n_heads = n_heads\n", + " self.d_head = d_model // n_heads\n", + "\n", + " self.q_proj = nn.Linear(d_in_q, d_model, bias=False)\n", + " self.k_proj = nn.Linear(d_in_kv, d_model, bias=False)\n", + " self.v_proj = nn.Linear(d_in_kv, d_model, bias=False)\n", + " self.out = nn.Linear(d_model, d_in_q, bias=False)\n", + "\n", + " def forward(self, x, y=None):\n", + " if y is None:\n", + " y = x\n", + "\n", + " B, N_q, _ = x.shape\n", + " N_kv = y.shape[1]\n", + "\n", + " q = self.q_proj(x).view(B, N_q, self.n_heads, self.d_head).transpose(1, 2)\n", + " k = self.k_proj(y).view(B, N_kv, self.n_heads, self.d_head).transpose(1, 2)\n", + " v = self.v_proj(y).view(B, N_kv, self.n_heads, self.d_head).transpose(1, 2)\n", + "\n", + " attn = torch.nn.functional.scaled_dot_product_attention(\n", + " q, k, v, dropout_p=0.0, is_causal=False\n", + " )\n", + "\n", + " attn = attn.transpose(1, 2).contiguous().view(B, N_q, -1)\n", + " return self.out(attn)\n", + "\n", + "\n", + "class Critic(nn.Module):\n", + " def __init__(\n", + " self,\n", + " inputs: int,\n", + " width_phi: int = 32,\n", + " depth_phi: int = 4,\n", + " tgt_encoded_dim: int = 16,\n", + " width_psi: int = 32,\n", + " depth_psi: int = 4,\n", + " n_tgts: int = 32,\n", + " obs_sat: int = 38,\n", + " dropout: float = 0.0,\n", + " ):\n", + " super(Critic, self).__init__()\n", + "\n", + " self.obs_sat = obs_sat\n", + " act_function = nn.ReLU\n", + "\n", + " self.n_tgts = n_tgts\n", + " # Define the number of features per target based on the input size and the number of targets\n", + " self.features_per_tgt = (inputs - self.obs_sat) // n_tgts\n", + "\n", + " input_size_phi = self.features_per_tgt\n", + "\n", + " layers_phi = []\n", + " layers_phi.append(nn.Linear(input_size_phi, width_phi))\n", + " layers_phi.append(act_function())\n", + " if dropout > 0:\n", + " layers_phi.append(nn.Dropout(dropout))\n", + " for _ in range(depth_phi - 1):\n", + " layers_phi.append(nn.Linear(width_phi, width_phi))\n", + " layers_phi.append(act_function())\n", + " if dropout > 0:\n", + " layers_phi.append(nn.Dropout(dropout))\n", + " layers_phi.append(nn.Linear(width_phi, tgt_encoded_dim))\n", + " self.phi = nn.Sequential(*layers_phi)\n", + "\n", + " input_size_psi = 2 * tgt_encoded_dim + self.obs_sat\n", + "\n", + " layers_psi = []\n", + " layers_psi.append(nn.Linear(input_size_psi, width_psi))\n", + " layers_psi.append(act_function())\n", + " if dropout > 0:\n", + " layers_psi.append(nn.Dropout(dropout))\n", + "\n", + " for _ in range(depth_psi - 1):\n", + " layers_psi.append(nn.Linear(width_psi, width_psi))\n", + " layers_psi.append(act_function())\n", + " if dropout > 0:\n", + " layers_psi.append(nn.Dropout(dropout))\n", + "\n", + " layers_psi.append(nn.Linear(width_psi, 1))\n", + " self.psi = nn.Sequential(*layers_psi)\n", + "\n", + " def forward(self, x):\n", + " if isinstance(x, dict) and \"obs\" in x:\n", + " x = x[\"obs\"]\n", + "\n", + " B = x.shape[0]\n", + " x_sat = x[:, : self.obs_sat]\n", + " x_tgts = x[:, self.obs_sat :]\n", + "\n", + " # Allows changes in the number of targets during runtime without changing internal variables as long as the input dimension is consistent with the number of targets\n", + " n_tgts = x_tgts.shape[1] // self.features_per_tgt\n", + " x_tgts = x_tgts.view(\n", + " B, n_tgts, self.features_per_tgt\n", + " ) # (B, n_tgts, features_per_tgt)\n", + "\n", + " latent_tgts = self.phi(x_tgts)\n", + "\n", + " latent = torch.cat(\n", + " [\n", + " x_sat,\n", + " torch.mean(latent_tgts, dim=1),\n", + " torch.max(latent_tgts, dim=1).values,\n", + " ],\n", + " dim=-1,\n", + " )\n", + "\n", + " critic_value = self.psi(latent).squeeze(-1) # (B,)\n", + "\n", + " return critic_value\n", + "\n", + "\n", + "class Actor(nn.Module):\n", + " def __init__(\n", + " self,\n", + " inputs: int,\n", + " width_phi: int = 32,\n", + " depth_phi: int = 4,\n", + " tgt_encoded_dim: int = 16,\n", + " num_heads: int = 2,\n", + " attention_dim: int = 32,\n", + " width_psi_img: int = 32,\n", + " depth_psi_img: int = 2,\n", + " n_tgts: int = 32,\n", + " obs_sat: int = 38,\n", + " non_imaging_actions: int = 1,\n", + " dropout: float = 0.0,\n", + " width_psi_sat: int = 32,\n", + " depth_psi_sat: int = 2,\n", + " width_phi_sat: int = 32,\n", + " depth_phi_sat: int = 2,\n", + " sat_attention_dim: int = 32,\n", + " sat_attention_heads: int = 2,\n", + " sat_encoded_dim: int = 32,\n", + " hierarchical: bool = False,\n", + " ):\n", + " super(Actor, self).__init__()\n", + "\n", + " self.obs_sat = obs_sat\n", + " act_function = nn.ReLU\n", + " self.n_tgts = n_tgts\n", + " # Define the number of features per target based on the input size and the number of targets\n", + " self.features_per_tgt = (inputs - self.obs_sat) // n_tgts\n", + " self.non_imaging_actions = non_imaging_actions\n", + " self.hierarchical = hierarchical\n", + "\n", + " self.phi = self._build_phi_model(\n", + " input_dim=self.features_per_tgt,\n", + " width=width_phi,\n", + " depth=depth_phi,\n", + " output_dim=tgt_encoded_dim,\n", + " act_function=act_function,\n", + " dropout=dropout,\n", + " )\n", + "\n", + " self.self_attention = MultiHeadSelfAttention(\n", + " tgt_encoded_dim, attention_dim, num_heads\n", + " )\n", + "\n", + " self.norm_self_attention = nn.LayerNorm(tgt_encoded_dim)\n", + "\n", + " self.psi_img = self._build_psi_model(\n", + " 3 * tgt_encoded_dim + sat_encoded_dim,\n", + " width=width_psi_img,\n", + " depth=depth_psi_img,\n", + " output_dim=tgt_encoded_dim,\n", + " act_function=act_function,\n", + " dropout=dropout,\n", + " add_norm=False,\n", + " )\n", + "\n", + " self.out_layer_psi_img = nn.Linear(tgt_encoded_dim, 1)\n", + "\n", + " self.phi_sat = self._build_phi_model(\n", + " input_dim=self.obs_sat,\n", + " width=width_phi_sat,\n", + " depth=depth_phi_sat,\n", + " output_dim=sat_encoded_dim,\n", + " act_function=act_function,\n", + " dropout=dropout,\n", + " )\n", + "\n", + " self.cross_attention = MultiHeadAttention(\n", + " sat_encoded_dim,\n", + " tgt_encoded_dim,\n", + " sat_attention_dim,\n", + " sat_attention_heads,\n", + " )\n", + "\n", + " self.norm_cross_attention = nn.LayerNorm(sat_encoded_dim)\n", + "\n", + " self.psi_sat = self._build_psi_model(\n", + " 2 * tgt_encoded_dim + sat_encoded_dim,\n", + " width=width_psi_sat,\n", + " depth=depth_psi_sat,\n", + " output_dim=sat_encoded_dim,\n", + " act_function=act_function,\n", + " dropout=dropout,\n", + " add_norm=False,\n", + " )\n", + "\n", + " out_dim_psi_sat = 1 + non_imaging_actions\n", + " self.out_layer_psi_sat = nn.Linear(sat_encoded_dim, out_dim_psi_sat)\n", + "\n", + " def _build_phi_model(\n", + " self,\n", + " input_dim: int,\n", + " width: int,\n", + " depth: int,\n", + " output_dim: int,\n", + " act_function,\n", + " dropout: float,\n", + " ):\n", + " layers = []\n", + " layers.append(nn.Linear(input_dim, width))\n", + " layers.append(act_function())\n", + " if dropout > 0:\n", + " layers.append(nn.Dropout(dropout))\n", + " for _ in range(depth - 1):\n", + " layers.append(nn.Linear(width, width))\n", + " layers.append(act_function())\n", + " if dropout > 0:\n", + " layers.append(nn.Dropout(dropout))\n", + "\n", + " layers.append(nn.Linear(width, output_dim))\n", + " return nn.Sequential(*layers)\n", + "\n", + " def _build_psi_model(\n", + " self,\n", + " input_dim: int,\n", + " width: int,\n", + " depth: int,\n", + " output_dim: int,\n", + " act_function,\n", + " dropout: float,\n", + " add_norm: bool,\n", + " ):\n", + " layers = []\n", + " layers.append(nn.Linear(input_dim, width))\n", + " layers.append(act_function())\n", + " if dropout > 0:\n", + " layers.append(nn.Dropout(dropout))\n", + " for _ in range(depth - 1):\n", + " layers.append(nn.Linear(width, width))\n", + " layers.append(act_function())\n", + " if dropout > 0:\n", + " layers.append(nn.Dropout(dropout))\n", + "\n", + " layers.append(nn.Linear(width, output_dim))\n", + " layers.append(act_function())\n", + "\n", + " return nn.Sequential(*layers)\n", + "\n", + " def forward(self, x):\n", + " if isinstance(x, dict) and \"obs\" in x:\n", + " x = x[\"obs\"]\n", + "\n", + " B = x.shape[0]\n", + "\n", + " x_sat = x[:, : self.obs_sat]\n", + " x_tgts = x[:, self.obs_sat :]\n", + "\n", + " # Allows changes in the number of targets during runtime without changing internal variables as long as the input dimension is consistent with the number of targets\n", + " n_tgts = x_tgts.shape[1] // self.features_per_tgt\n", + " x_tgts = x_tgts.view(\n", + " B, n_tgts, self.features_per_tgt\n", + " ) # (B, n_tgts, features_per_tgt)\n", + " latent_tgts = self.phi(x_tgts) # (B, n_tgts, tgt_encoded_dim)\n", + " latent_sat = self.phi_sat(x_sat) # (B, sat_encoded_dim)\n", + "\n", + " attention_out_self = self.self_attention(\n", + " latent_tgts\n", + " ) # (B, n_tgts, tgt_encoded_dim)\n", + "\n", + " latent_tgts_self = self.norm_self_attention(latent_tgts + attention_out_self)\n", + "\n", + " latent_tgts = self.psi_img(\n", + " torch.cat(\n", + " [\n", + " latent_tgts_self,\n", + " torch.mean(latent_tgts_self, dim=1)\n", + " .unsqueeze(1)\n", + " .expand(-1, n_tgts, -1),\n", + " torch.max(latent_tgts_self, dim=1)\n", + " .values.unsqueeze(1)\n", + " .expand(-1, n_tgts, -1),\n", + " latent_sat.unsqueeze(1).expand(-1, n_tgts, -1),\n", + " ],\n", + " dim=-1,\n", + " )\n", + " ) # (B, n_tgts, tgt_encoded_dim\n", + "\n", + " # Cross attention with satellite features\n", + " sat_attention_out = self.cross_attention(\n", + " latent_sat.unsqueeze(1), latent_tgts\n", + " ) # (B, 1, sat_attention_dim)\n", + "\n", + " latent_sat = self.norm_cross_attention(\n", + " latent_sat + sat_attention_out.squeeze(1)\n", + " ) # (B, sat_encoded_dim)\n", + "\n", + " latent_sat = self.psi_sat(\n", + " torch.cat(\n", + " [\n", + " latent_sat,\n", + " torch.mean(latent_tgts, dim=1),\n", + " torch.max(latent_tgts, dim=1).values,\n", + " ],\n", + " dim=-1,\n", + " )\n", + " )\n", + "\n", + " logits_tgts = self.out_layer_psi_img(latent_tgts).squeeze(-1)\n", + "\n", + " psi_sat_out = self.out_layer_psi_sat(latent_sat) # (B, non_imaging_actions + 1)\n", + "\n", + " img_modulated = psi_sat_out[:, 0:1]\n", + " non_img_logit = psi_sat_out[:, 1:]\n", + "\n", + " if self.hierarchical:\n", + " return torch.cat([non_img_logit, img_modulated, logits_tgts], dim=1)\n", + "\n", + " return torch.cat([non_img_logit, logits_tgts + img_modulated], dim=1)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A custom `RLModule` is created, inheriting from the existing `PPOTorchRLModule`. This new module instantiates the new actor and critic modules and re-defines the pi head." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class CustomModule(PPOTorchRLModule, nn.Module):\n", + " def setup(self):\n", + " catalog = self.config.get_catalog()\n", + " is_stateful = isinstance(\n", + " catalog.actor_critic_encoder_config.base_encoder_config,\n", + " RecurrentEncoderConfig,\n", + " )\n", + " if is_stateful:\n", + " self.config.inference_only = False\n", + "\n", + " if self.config.inference_only and self.framework == \"torch\":\n", + " catalog.actor_critic_encoder_config.inference_only = True\n", + "\n", + " self.encoder = lambda x: {ENCODER_OUT: {ACTOR: x, CRITIC: x}}\n", + "\n", + " config_dict = self.config.model_config_dict\n", + "\n", + " self.pi_head = Actor(\n", + " inputs=self.config.observation_space.shape[0],\n", + " n_tgts=config_dict[\"n_targets\"],\n", + " obs_sat=config_dict[\"obs_sat\"],\n", + " width_phi=config_dict[\"width_phi\"],\n", + " depth_phi=config_dict[\"depth_phi\"],\n", + " tgt_encoded_dim=config_dict[\"tgt_encoded_dim\"],\n", + " num_heads=config_dict[\"num_heads\"],\n", + " attention_dim=config_dict[\"attention_dim\"],\n", + " width_psi_img=config_dict[\"width_psi_img\"],\n", + " depth_psi_img=config_dict[\"depth_psi_img\"],\n", + " dropout=config_dict.get(\"dropout\", 0.0),\n", + " non_imaging_actions=config_dict.get(\"non_imaging_actions\", 1),\n", + " width_psi_sat=config_dict[\"width_psi_sat\"],\n", + " depth_psi_sat=config_dict[\"depth_psi_sat\"],\n", + " sat_attention_dim=config_dict[\"sat_attention_dim\"],\n", + " sat_attention_heads=config_dict[\"sat_attention_heads\"],\n", + " width_phi_sat=config_dict[\"width_phi_sat\"],\n", + " depth_phi_sat=config_dict[\"depth_phi_sat\"],\n", + " sat_encoded_dim=config_dict[\"sat_encoded_dim\"],\n", + " )\n", + "\n", + " if not self.config.inference_only or self.framework != \"torch\":\n", + " self.vf = Critic(\n", + " inputs=self.config.observation_space.shape[0],\n", + " n_tgts=config_dict[\"n_targets\"],\n", + " obs_sat=config_dict[\"obs_sat\"],\n", + " width_phi=config_dict[\"critic_width_phi\"],\n", + " depth_phi=config_dict[\"critic_depth_phi\"],\n", + " tgt_encoded_dim=config_dict[\"critic_tgt_encoded_dim\"],\n", + " width_psi=config_dict[\"critic_width_psi\"],\n", + " depth_psi=config_dict[\"critic_depth_psi\"],\n", + " dropout=config_dict.get(\"dropout\", 0.1),\n", + " )\n", + " self._inference_only_state_dict_keys = {}\n", + "\n", + " self.action_dist_cls = catalog.get_action_dist_cls(framework=self.framework)\n", + "\n", + " def pi(\n", + " self, batch: Dict[str, TensorType], inference: bool = False\n", + " ) -> Dict[str, TensorType]:\n", + " pi_outs = {}\n", + "\n", + " logits = self.pi_head(batch)\n", + " discrete_action_dist = TorchCategorical.from_logits(logits)\n", + "\n", + " if inference:\n", + " discrete_action = discrete_action_dist.to_deterministic().sample()\n", + " else:\n", + " discrete_action = discrete_action_dist.sample()\n", + "\n", + " discrete_action_logp = discrete_action_dist.logp(discrete_action)\n", + "\n", + " pi_outs[Columns.ACTION_LOGP] = discrete_action_logp\n", + " pi_outs[Columns.ACTION_DIST_INPUTS] = logits\n", + "\n", + " pi_outs[Columns.ACTIONS] = discrete_action\n", + " return pi_outs\n", + "\n", + " @override(TorchRLModule)\n", + " def _forward_inference(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:\n", + " return self.pi(batch, inference=True)\n", + "\n", + " @override(TorchRLModule)\n", + " def _forward_exploration(\n", + " self, batch: Dict[str, TensorType], **kwargs\n", + " ) -> Dict[str, TensorType]:\n", + " return self.pi(batch, inference=False)\n", + "\n", + " @override(TorchRLModule)\n", + " def _forward_train(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:\n", + " outs = {}\n", + " outs.update(self.pi(batch))\n", + " vf_out = self.vf(batch)\n", + " outs[Columns.VF_PREDS] = vf_out.squeeze(-1)\n", + " return outs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Then, the training hyperparameters and modules parameters are defined, as well as the custom RLModule." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from ray.rllib.core.rl_module.rl_module import RLModuleSpec\n", + "\n", + "training_args = dict(\n", + " lr=[\n", + " [0, 0.00033003435881682255],\n", + " [40000, 0.00033003435881682255 / 16.749479444886223],\n", + " ],\n", + " gamma=0.999,\n", + " train_batch_size=int(300 * 10 * 3),\n", + " num_sgd_iter=30,\n", + " lambda_=0.8713548569911232,\n", + " use_kl_loss=False,\n", + " clip_param=0.14701727973480344,\n", + " grad_clip=0.3104924935285628,\n", + " vf_clip_param=2.0,\n", + " entropy_coeff=[[0, 0.023694512589767867], [750_000, 0.0]],\n", + ")\n", + "\n", + "rl_module_args = dict(\n", + " model_config_dict={\n", + " \"n_targets\": 40,\n", + " \"obs_sat\": 38,\n", + " \"width_phi\": 256,\n", + " \"depth_phi\": 2,\n", + " \"width_psi_img\": 128,\n", + " \"depth_psi_img\": 4,\n", + " \"tgt_encoded_dim\": 128,\n", + " \"attention_depth\": 1,\n", + " \"num_heads\": 2,\n", + " \"attention_dim\": 128,\n", + " \"width_phi_sat\": 256,\n", + " \"depth_phi_sat\": 2,\n", + " \"width_psi_sat\": 128,\n", + " \"depth_psi_sat\": 4,\n", + " \"sat_attention_dim\": 128,\n", + " \"sat_attention_heads\": 2,\n", + " \"sat_encoded_dim\": 128,\n", + " \"act_function\": \"ReLU\",\n", + " \"critic_tgt_encoded_dim\": 128,\n", + " \"critic_width_phi\": 256,\n", + " \"critic_depth_phi\": 2,\n", + " \"critic_width_psi\": 64,\n", + " \"critic_depth_psi\": 3,\n", + " \"dropout\": 0.1,\n", + " \"critic_dropout\": 0.1,\n", + " \"non_imaging_actions\": 1,\n", + " },\n", + " rl_module_spec=RLModuleSpec(module_class=CustomModule),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from bsk_rl.utils.rllib.callbacks import WrappedEpisodeDataCallbacks\n", + "from bsk_rl.utils.rllib.discounting import TimeDiscountedGAEPPOTorchLearner\n", + "from ray.rllib.algorithms.ppo import PPOConfig\n", + "\n", + "N_CPUS = 12\n", + "\n", + "\n", + "def example_data_callback(env):\n", + " reward = env.rewarder.cum_reward\n", + " reward = sum(reward.values())\n", + " orbits = env.simulator.sim_time / (95 * 60)\n", + "\n", + " data = dict(\n", + " reward=reward,\n", + " alive=float(env.satellite.is_alive()),\n", + " battery_status_valid=float(env.satellite.dynamics.battery_valid()),\n", + " orbits_complete=orbits,\n", + " )\n", + " if orbits > 0:\n", + " data[\"reward_per_orbit\"] = reward / orbits\n", + " if not env.satellite.is_alive():\n", + " data[\"orbits_complete_partial_only\"] = orbits\n", + "\n", + " return data\n", + "\n", + "\n", + "env_args_training = dict(\n", + " satellite=CS_F_Bayesian(\"C_F\", sat_args, belief_update_func=belief_update_func),\n", + " scenario=scenario,\n", + " rewarder=rewarder,\n", + " sat_arg_randomizer=sat_arg_randomizer,\n", + " sim_rate=0.5,\n", + " max_step_duration=300.0,\n", + " time_limit=95 * 60 * 3, # Three orbits\n", + " log_level=\"INFO\",\n", + " failure_penalty=0.0,\n", + ")\n", + "\n", + "ppo_config = (\n", + " PPOConfig()\n", + " .training(\n", + " **training_args,\n", + " learner_class=TimeDiscountedGAEPPOTorchLearner,\n", + " )\n", + " .env_runners(num_env_runners=N_CPUS - 1, sample_timeout_s=1000.0)\n", + " .environment(\n", + " env=\"SatelliteTasking-RLlib\",\n", + " env_config=dict(\n", + " **env_args_training, episode_data_callback=example_data_callback\n", + " ),\n", + " )\n", + " .reporting(\n", + " metrics_num_episodes_for_smoothing=1,\n", + " metrics_episode_collection_timeout_s=180,\n", + " )\n", + " .checkpointing(export_native_model_files=True)\n", + " .framework(framework=\"torch\")\n", + " .api_stack(\n", + " enable_rl_module_and_learner=True,\n", + " enable_env_runner_and_connector_v2=True,\n", + " )\n", + " .callbacks(WrappedEpisodeDataCallbacks)\n", + ")\n", + "ppo_config.rl_module(**rl_module_args)\n", + "\n", + "# Uncomment to run training\n", + "# ray.init(\n", + "# ignore_reinit_error=True,\n", + "# num_cpus=N_CPUS,\n", + "# object_store_memory=2_000_000_000, # 2 GB\n", + "# )\n", + "# results = tune.run(\n", + "# \"PPO\",\n", + "# config=ppo_config.to_dict(),\n", + "# stop={\n", + "# \"num_env_steps_sampled_lifetime\": 264\n", + "# }, # Total number of steps to train the model. Originally 774,000\n", + "# checkpoint_freq=1,\n", + "# checkpoint_at_end=True,\n", + "# )\n", + "\n", + "# ray.shutdown()" ] } ], "metadata": { "kernelspec": { - "display_name": ".venv_update_cloud_env_JAIS", + "display_name": ".venv_update_cloud_ex (3.11.9.final.0)", "language": "python", "name": "python3" },