[#144] Resource reward randomization

Author: Mark2000

docs/source/release_notes.rst

@@ -35,6 +35,8 @@ Development - |version|
 * Add properties in spacecraft dynamics for orbital element observations.
 * Fix an issue with failure penalties in the PettingZoo environment when the rewarder
   does not return a reward for a satellite.
+* Allow for per-episode randomization of :class:`ResourceReward` weights and observation
+  of those weights with :class:`ResourceRewardWeight`.
 
 
 Version 1.1.0

src/bsk_rl/act/actions.py

@@ -4,7 +4,7 @@
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any
 
-from bsk_rl.utils.functional import AbstractClassProperty
+from bsk_rl.utils.functional import AbstractClassProperty, Resetable
 
 if TYPE_CHECKING:  # pragma: no cover
     from gymnasium import spaces
@@ -29,7 +29,7 @@ def select_action_builder(satellite: "Satellite") -> "ActionBuilder":
         raise NotImplementedError("Heterogenous action builders not supported.")
 
 
-class ActionBuilder(ABC):
+class ActionBuilder(ABC, Resetable):
 
     def __init__(self, satellite: "Satellite") -> None:
         """Base class for all action builders.
@@ -43,6 +43,21 @@ def __init__(self, satellite: "Satellite") -> None:
         for act in self.action_spec:
             act.link_satellite(self.satellite)
 
+    def reset_overwrite_previous(self) -> None:
+        """Perform any once-per-episode setup."""
+        for act in self.action_spec:
+            act.reset_overwrite_previous()
+
+    def reset_pre_sim_init(self) -> None:
+        """Perform any once-per-episode setup."""
+        for act in self.action_spec:
+            act.reset_pre_sim_init()
+
+    def reset_during_sim_init(self) -> None:
+        """Perform any once-per-episode setup."""
+        for act in self.action_spec:
+            act.reset_during_sim_init()
+
     def reset_post_sim_init(self) -> None:
         """Perform any once-per-episode setup."""
         self.simulator = self.satellite.simulator  # already a proxy
@@ -68,7 +83,7 @@ def set_action(self, action: Any) -> None:
         pass
 
 
-class Action(ABC):
+class Action(ABC, Resetable):
     builder_type: type[ActionBuilder] = AbstractClassProperty()  #: :meta private:
 
     def __init__(self, name: str = "act") -> None:
@@ -101,10 +116,6 @@ def link_simulator(self, simulator: "Simulator") -> None:
         """
         self.simulator = simulator  # already a proxy
 
-    def reset_post_sim_init(self) -> None:  # pragma: no cover
-        """Perform any once-per-episode setup."""
-        pass
-
     @abstractmethod
     def set_action(self, action: Any) -> None:  # pragma: no cover
         """Execute code to perform an action."""

src/bsk_rl/data/resource_data.py

@@ -5,9 +5,10 @@
 """
 
 import logging
-from typing import TYPE_CHECKING, Callable
+from typing import Callable, Union
 
 from bsk_rl.data.base import Data, DataStore, GlobalReward
+from bsk_rl.obs import ResourceRewardWeight
 
 logger = logging.getLogger(__name__)
 
@@ -57,7 +58,9 @@ class ResourceReward(GlobalReward):
     data_store_type = ResourceDataStore
 
     def __init__(
-        self, reward_weight: float = 1.0, resource_fn: Callable = lambda sat: 0.0
+        self,
+        reward_weight: Union[float, Callable] = 1.0,
+        resource_fn: Callable = lambda sat: 0.0,
     ) -> None:
         """Rewards for an arbitrary resource.
 
@@ -76,17 +79,33 @@ def __init__(
             resource_fn = lambda sat: sat.simulator.sim_time
             reward_weight = -1e-3  # is negative, because time increases over time
 
-
         Args:
             reward_weight: [reward/resource] Scaling factor to apply to changes in resource
-                level to yield reward.
+                level to yield reward. Can be a float or a function that randomizes the reward
+                weight per-episode.
             resource_fn: Function to call to get the resource level for each satellite.
         """
         super().__init__()
-        self.reward_weight = reward_weight
+        self._reward_weight = reward_weight
         self.resource_fn = resource_fn
         self.data_store_kwargs = dict(resource_fn=resource_fn)
 
+    def reset_pre_sim_init(self) -> None:
+        """Reset the reward weight before simulation initialization."""
+        if callable(self._reward_weight):
+            self.reward_weight = self._reward_weight()
+        else:
+            self.reward_weight = self._reward_weight
+        return super().reset_pre_sim_init()
+
+    def reset_post_sim_init(self) -> None:
+        """Add the reward weight to each satellite's observation spec."""
+        for satellite in self.scenario.satellites:
+            for obs in satellite.observation_builder.observation_spec:
+                if isinstance(obs, ResourceRewardWeight):
+                    obs.weight_vector.append(self.reward_weight)
+        return super().reset_post_sim_init()
+
     def calculate_reward(
         self, new_data_dict: dict[str, ResourceData]
     ) -> dict[str, float]:

src/bsk_rl/obs/__init__.py

@@ -32,12 +32,14 @@ class MyObservationSatellite(Satellite):
 * :class:`Time` - Add simulation time to the observation.
 * :class:`OpportunityProperties` - Add information about upcoming targets or other ground access points to the observation.
 * :class:`Eclipse` - Add a tuple of the next orbit start and end.
+* :class:`ResourceRewardWeight` - Reports the weights of any randomized :class:`ResourceReward`.
 """
 
 from bsk_rl.obs.observations import (
     Eclipse,
     Observation,
     OpportunityProperties,
+    ResourceRewardWeight,
     SatProperties,
     Time,
 )
@@ -51,4 +53,5 @@ class MyObservationSatellite(Satellite):
     "Time",
     "OpportunityProperties",
     "Eclipse",
+    "ResourceRewardWeight",
 ]

src/bsk_rl/obs/observations.py

@@ -3,13 +3,12 @@
 import logging
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Union
 
 import numpy as np
-from Basilisk.utilities import orbitalMotion
 from gymnasium import spaces
 
-from bsk_rl.utils.functional import vectorize_nested_dict
+from bsk_rl.utils.functional import Resetable, vectorize_nested_dict
 from bsk_rl.utils.orbital import rv2HN
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -49,7 +48,7 @@ def nested_obs_to_space(obs_dict, dtype):
         raise TypeError(f"Cannot convert {obs_dict} to gym space.")
 
 
-class ObservationBuilder:
+class ObservationBuilder(Resetable):
     def __init__(
         self,
         satellite: "Satellite",
@@ -80,6 +79,21 @@ def __init__(
                 name_counts[obs.name] = 1
             obs.link_satellite(self.satellite)
 
+    def reset_overwrite_previous(self) -> None:
+        """Perform any once-per-episode setup."""
+        for obs in self.observation_spec:
+            obs.reset_overwrite_previous()
+
+    def reset_pre_sim_init(self) -> None:
+        """Perform any once-per-episode setup."""
+        for obs in self.observation_spec:
+            obs.reset_pre_sim_init()
+
+    def reset_during_sim_init(self) -> None:
+        """Perform any once-per-episode setup."""
+        for obs in self.observation_spec:
+            obs.reset_during_sim_init()
+
     def reset_post_sim_init(self) -> None:
         """Perform any once-per-episode setup."""
         self.simulator = self.satellite.simulator  # already a proxy
@@ -143,7 +157,7 @@ def observation_description(self) -> Any:
         return self.obs_array_keys()
 
 
-class Observation(ABC):
+class Observation(ABC, Resetable):
     """Base observations class."""
 
     def __init__(self, name: str = "obs") -> None:
@@ -176,10 +190,6 @@ def link_simulator(self, simulator: "Simulator") -> None:
         """
         self.simulator = simulator  # already a proxy
 
-    def reset_post_sim_init(self) -> None:  # pragma: no cover
-        """Perform any once-per-episode setup."""
-        pass
-
     @abstractmethod  # pragma: no cover
     def get_obs(self) -> Any:
         """Return the observation."""
@@ -481,5 +491,30 @@ def get_obs(self):
         ]
 
 
+class ResourceRewardWeight(Observation):
+    """Observation for the weight of any :class:`ResourceReward`."""
+
+    def __init__(
+        self, name="resource_reward_weight", norm: Union[float, np.ndarray] = 1.0
+    ):
+        """Include the resource reward weight in the observation.
+
+        Args:
+            name: Name of the observation.
+            norm: Value to normalize the resource reward weight by. If a vector, it should
+                be the same length as the number of resource rewards in the environment.
+        """
+        super().__init__(name=name)
+        self.norm = norm
+
+    def reset_overwrite_previous(self) -> None:  # pragma: no cover
+        """Prepare weights to be saved by rewarders."""
+        self.weight_vector = []
+
+    def get_obs(self) -> float:
+        """Return the resource reward weight."""
+        return np.array(self.weight_vector) / self.norm
+
+
 __doc_title__ = "Backend"
 __all__ = ["ObservationBuilder"]

src/bsk_rl/sats/satellite.py

@@ -142,6 +142,8 @@ def reset_overwrite_previous(self) -> None:
         self._timed_terminal_event_name = None
         self._is_alive = True
         self.time_of_death = None
+        self.observation_builder.reset_overwrite_previous()
+        self.action_builder.reset_overwrite_previous()
 
     @vizard.visualize
     def create_vizard_data(self, color, vizSupport=None) -> None:
@@ -160,6 +162,8 @@ def reset_pre_sim_init(self) -> None:
             oe=self.sat_args["oe"],
             mu=self.sat_args["mu"],
         )
+        self.observation_builder.reset_pre_sim_init()
+        self.action_builder.reset_pre_sim_init()
 
     def set_simulator(self, simulator: "Simulator"):
         """Set the simulator for models.
@@ -203,6 +207,12 @@ def set_fsw(self, fsw_rate: float) -> "fsw.FSWModel":
         self.fsw = proxy(fsw)
         return fsw
 
+    def reset_during_sim_init(self) -> None:
+        """Called during environment reset, during Basilisk simulation initialization."""
+        self.observation_builder.reset_during_sim_init()
+        self.action_builder.reset_during_sim_init()
+        return super().reset_during_sim_init()
+
     def reset_post_sim_init(self) -> None:
         """Called during environment reset, after Basilisk simulation initialization."""
         self.observation_builder.reset_post_sim_init()

tests/integration/obs/test_int_observations.py

@@ -212,3 +212,37 @@ class GroundSat(sats.AccessSatellite):
     def test_ground_station_state(self):
         observation, info = self.env.reset()
         assert sum(observation) > 0  # Check that there are downlink opportunities
+
+
+class TestResourceRewardWeight:
+    class ResourceSat(sats.ImagingSatellite):
+        dyn_type = dyn.ImagingDynModel
+        fsw_type = fsw.ImagingFSWModel
+        observation_spec = [obs.ResourceRewardWeight()]
+        action_spec = [act.Drift()]
+
+    env = gym.make(
+        "SatelliteTasking-v1",
+        satellite=ResourceSat(
+            "ResourceSat",
+            obs_type=dict,
+        ),
+        scenario=UniformTargets(n_targets=0),
+        rewarder=(
+            data.ResourceReward(
+                resource_fn=lambda sat: sat.simulator.sim_time, reward_weight=0.1
+            ),
+            data.ResourceReward(
+                resource_fn=lambda sat: sat.simulator.sim_time, reward_weight=1.0
+            ),
+        ),
+        sim_rate=1.0,
+        max_step_duration=10.0,
+        disable_env_checker=True,
+    )
+
+    def test_resource_reward_weight(self):
+        observation, info = self.env.reset()
+        observation, reward, terminated, truncated, info = self.env.step(0)
+        assert reward == 1.0 + 10.0
+        assert (observation["resource_reward_weight"] == [0.1, 1.0]).all()

tests/unittest/data/test_data.py

@@ -339,6 +339,20 @@ def test_compare_log_states(self):
 class TestResourceReward:
     def test_calculate_reward(self):
         dm = ResourceReward(reward_weight=2.0)
+        dm.reset_overwrite_previous()
+        dm.reset_pre_sim_init()
+        reward = dm.calculate_reward(
+            {
+                "sat1": ResourceData(1.0),
+                "sat2": ResourceData(-2.0),
+            }
+        )
+        assert reward == {"sat1": 2.0, "sat2": -4.0}
+
+    def test_calculate_random_reward(self):
+        dm = ResourceReward(reward_weight=lambda: 2.0)
+        dm.reset_overwrite_previous()
+        dm.reset_pre_sim_init()
         reward = dm.calculate_reward(
             {
                 "sat1": ResourceData(1.0),
@@ -350,6 +364,7 @@ def test_calculate_reward(self):
     def test_read_reward(self):
         dm = ResourceReward(resource_fn=lambda sat: sat.resource_level)
         dm.reset_overwrite_previous()
+        dm.reset_pre_sim_init()
         sat = MagicMock()
         dm.create_data_store(sat)
         sat.resource_level = 3.0

tests/unittest/obs/test_observations.py

@@ -221,3 +221,11 @@ def test_obs(self):
         ob.satellite = MagicMock()
         ob.satellite.trajectory.next_eclipse.return_value = (20.0, 30.0)
         assert ob.get_obs() == [0.1, 0.2]
+
+
+class TestResourceRewardWeight:
+    def test_obs(self):
+        ob = obs.ResourceRewardWeight()
+        ob.reset_overwrite_previous()
+        ob.weight_vector.append(1.0)
+        assert ob.get_obs() == [1.0]

tests/unittest/sats/test_access_satellite.py

@@ -310,6 +310,8 @@ def make_sat(self):
     @patch("bsk_rl.sats.Satellite.reset_pre_sim_init")
     def test_reset_pre_sim_init(self, mock_reset):
         sat = self.make_sat()
+        sat.observation_builder = MagicMock()
+        sat.action_builder = MagicMock()
         sat.reset_overwrite_previous()
         targets = [MagicMock()] * 5
         for target in targets: