[Issue #370] End downlink actions when the buffer empties

docs/source/release_notes.rst

@@ -5,6 +5,11 @@ Development - |version|
 -----------------------
 .. *Release Date: MMM. DD, YYYY*
 
+* Add opt-in buffer-conditioned termination to :class:`~bsk_rl.act.Downlink` with
+  ``variable_duration_downlink=True``. Downlink actions stop once storage is empty,
+  subject to a configurable 10-second minimum when the buffer is empty at action
+  start. A shorter configured action duration remains the cap. Existing fixed-duration
+  behavior remains the default. See issue #370.
 * Add optional hybrid action spaces with ``hybrid=True`` on supported discrete
   actions. Mode, broadcast, and ground-imaging actions accept a requested duration.
   Ground imaging instead uses a desired capture-time offset from the current

src/bsk_rl/act/actions.py

@@ -116,6 +116,10 @@ def link_simulator(self, simulator: "Simulator") -> None:
         """
         self.simulator = simulator  # already a proxy
 
+    def _clear_action_events(self) -> None:
+        """Remove simulator events that should expire on the next action selection."""
+        pass
+
     @abstractmethod
     def set_action(self, action: Any) -> None:  # pragma: no cover
         """Execute code to perform an action."""

src/bsk_rl/act/discrete_actions.py

@@ -101,6 +101,8 @@ def set_action(self, action: Union[int, tuple[int, float]]) -> None:
         only their local index, even when the overall action space is hybrid.
         """
         action_index, parameter = clean_action(action, hybrid=self.hybrid)
+        for act in self.action_spec:
+            act._clear_action_events()
         self.satellite.disable_timed_terminal_event()
         if not np.issubdtype(type(action_index), np.integer):
             logger.warning(
@@ -326,6 +328,9 @@ def __init__(
         name: Optional[str] = None,
         duration: Optional[float] = None,
         hybrid: bool = False,
+        variable_duration_downlink: bool = False,
+        empty_storage_threshold_bits: float = 1.0,
+        minimum_duration: float = 10.0,
     ):
         """Action to transmit data from the data buffer (:class:`~bsk_rl.sim.fsw.ImagingFSWModel.action_downlink`).
 
@@ -335,12 +340,100 @@ def __init__(
 
         Args:
             name: Action name.
-            duration: Time to task action, in seconds.
+            duration: Maximum time to task action, in seconds.
             hybrid: Whether to use a hybrid action space.
+            variable_duration_downlink: If true, stop early once storage is empty.
+            empty_storage_threshold_bits: Storage level considered empty, in bits.
+            minimum_duration: Minimum downlink time in seconds, including when
+                the buffer is empty at action start. The configured action
+                duration remains a hard maximum if it is shorter.
         """
         super().__init__(
             fsw_action="action_downlink", name=name, duration=duration, hybrid=hybrid
         )
+        if (
+            not np.isfinite(empty_storage_threshold_bits)
+            or empty_storage_threshold_bits < 0
+        ):
+            raise ValueError(
+                "empty_storage_threshold_bits must be finite and non-negative."
+            )
+        if not np.isfinite(minimum_duration) or minimum_duration < 0:
+            raise ValueError("minimum_duration must be finite and non-negative.")
+        self.variable_duration_downlink = bool(variable_duration_downlink)
+        self.empty_storage_threshold_bits = float(empty_storage_threshold_bits)
+        self.minimum_duration = float(minimum_duration)
+        self._downlink_empty_event_name = None
+        self._downlink_earliest_stop_time = None
+
+    def _storage_is_empty(self) -> bool:
+        """Return whether the satellite data buffer is effectively empty."""
+        return (
+            float(self.satellite.dynamics.storage_level)
+            <= self.empty_storage_threshold_bits
+        )
+
+    def reset_overwrite_previous(self) -> None:
+        """Forget event state from the previous episode."""
+        self._downlink_empty_event_name = None
+        self._downlink_earliest_stop_time = None
+
+    def _clear_action_events(self) -> None:
+        """Delete the active empty-buffer event, if one exists."""
+        if (
+            self._downlink_empty_event_name is not None
+            and self._downlink_empty_event_name in self.simulator.eventMap
+        ):
+            self.simulator.delete_event(self._downlink_empty_event_name)
+        self._downlink_empty_event_name = None
+        self._downlink_earliest_stop_time = None
+
+    def _downlink_empty_event_ready(self) -> bool:
+        """Wait for the configured minimum duration and an empty buffer."""
+        earliest_stop = self._downlink_earliest_stop_time
+        return (
+            earliest_stop is not None
+            and float(self.simulator.sim_time) >= earliest_stop
+            and self._storage_is_empty()
+        )
+
+    def _finish_downlink_early(self) -> None:
+        """Request a new action and keep event state available for cleanup."""
+        self.satellite.logger.info("Downlink storage empty")
+        self.satellite.requires_retasking = True
+        self._downlink_earliest_stop_time = None
+
+    def _enable_downlink_empty_event(self) -> None:
+        """Stop the action when its storage buffer is empty."""
+        self._clear_action_events()
+        # Only require the guard when there is no data to transmit at action start.
+        # If a nonempty buffer drains, end as soon as the next FSW event observes it.
+        minimum_step = (
+            max(self.minimum_duration, float(self.simulator.sim_rate))
+            if self._storage_is_empty()
+            else 0.0
+        )
+        self._downlink_earliest_stop_time = (
+            float(self.simulator.sim_time) + minimum_step
+        )
+        self._downlink_empty_event_name = f"downlink_empty_{self.satellite.name}"
+        self.simulator.createNewEvent(
+            self._downlink_empty_event_name,
+            mc.sec2nano(self.satellite.fsw.fsw_rate),
+            True,
+            conditionFunction=lambda sim: self._downlink_empty_event_ready(),
+            actionFunction=lambda sim: self._finish_downlink_early(),
+            terminal=self.satellite.variable_interval,
+        )
+
+    def set_action(
+        self, action: Union[int, tuple[int, float]], prev_action_key=None
+    ) -> str:
+        """Activate downlink and optionally stop when storage is empty."""
+        action_key = super().set_action(action, prev_action_key=prev_action_key)
+        if self.variable_duration_downlink:
+            self._enable_downlink_empty_event()
+        return action_key
 
 
 class Scan(DiscreteFSWAction):

tests/integration/act/test_int_actions_discrete.py

@@ -146,6 +146,65 @@ def test_image_by_name(self):
         self.env.step(target.id)
 
 
+class TestVariableDurationDownlink:
+    class ImageSat(sats.ImagingSatellite):
+        dyn_type = dyn.GroundStationDynModel
+        fsw_type = fsw.ImagingFSWModel
+        observation_spec: ClassVar[list[obs.Observation]] = [obs.Time()]
+        action_spec: ClassVar[list[act.Action]] = [
+            act.Downlink(variable_duration_downlink=True),
+            act.Drift(duration=20.0),
+        ]
+
+    env = gym.make(
+        "SatelliteTasking-v1",
+        satellite=ImageSat(
+            "EO_empty_buffer",
+            sat_args=ImageSat.default_sat_args(
+                oe=random_orbit,
+                instrumentBaudRate=1.0,
+                dataStorageCapacity=3.0,
+                transmitterBaudRate=-1.0,
+            ),
+        ),
+        scenario=UniformTargets(n_targets=0),
+        rewarder=data.UniqueImageReward(),
+        sim_rate=1.0,
+        time_limit=100.0,
+        max_step_duration=100.0,
+        disable_env_checker=True,
+    )
+
+    def test_empty_buffer_action_waits_for_minimum_then_retasks(self):
+        self.env.reset(seed=0)
+        start_time = self.env.unwrapped.simulator.sim_time
+
+        self.env.step(0)
+
+        elapsed = self.env.unwrapped.simulator.sim_time - start_time
+        assert elapsed >= 10.0
+        assert elapsed < 100.0
+
+
+    def test_downlink_event_is_removed_on_retask_and_reset(self):
+        self.env.reset(seed=0)
+        simulator = self.env.unwrapped.simulator
+        event_name = "downlink_empty_EO_empty_buffer"
+
+        self.env.step(0)
+        assert event_name in simulator.eventMap
+        self.env.step(1)
+        assert event_name not in simulator.eventMap
+
+        self.env.step(0)
+        assert event_name in simulator.eventMap
+        del simulator
+        self.env.reset(seed=0)
+        assert event_name not in self.env.unwrapped.simulator.eventMap
+        self.env.step(0)
+        assert self.env.unwrapped.simulator.sim_time == approx(10.0)
+
+
 ###########################
 # Individual Action Tests #
 ###########################

tests/unittest/act/test_actions.py

@@ -133,6 +133,66 @@ def test_set_action_reset(self):
         fswact.satellite.fsw.action_fsw.assert_called_once()
 
 
+class TestVariableDurationDownlink:
+    @staticmethod
+    def make_downlink(storage_level, *, minimum_duration=10.0):
+        downlink = act.Downlink(
+            variable_duration_downlink=True,
+            minimum_duration=minimum_duration,
+            empty_storage_threshold_bits=1.0,
+        )
+        downlink.satellite = MagicMock()
+        downlink.satellite.name = "sat"
+        downlink.satellite.dynamics.storage_level = storage_level
+        downlink.satellite.fsw.fsw_rate = 2.0
+        downlink.satellite.variable_interval = True
+        downlink.simulator = SimpleNamespace(
+            sim_rate=1.0,
+            sim_time=20.0,
+            eventMap={},
+            createNewEvent=MagicMock(),
+            delete_event=MagicMock(),
+        )
+        return downlink
+
+    def test_empty_at_start_waits_for_minimum_duration(self):
+        downlink = self.make_downlink(storage_level=0.0)
+
+        downlink._enable_downlink_empty_event()
+
+        assert downlink._downlink_earliest_stop_time == 30.0
+        assert not downlink._downlink_empty_event_ready()
+        downlink.simulator.sim_time = 30.0
+        assert downlink._downlink_empty_event_ready()
+
+    def test_nonempty_buffer_stops_as_soon_as_it_drains(self):
+        downlink = self.make_downlink(storage_level=5.0)
+
+        downlink._enable_downlink_empty_event()
+
+        assert downlink._downlink_earliest_stop_time == 20.0
+        assert not downlink._downlink_empty_event_ready()
+        downlink.satellite.dynamics.storage_level = 0.5
+        assert downlink._downlink_empty_event_ready()
+
+    def test_finishing_and_disabling_event_cleans_up_state(self):
+        downlink = self.make_downlink(storage_level=0.0)
+        downlink._enable_downlink_empty_event()
+        event_name = downlink._downlink_empty_event_name
+        downlink.simulator.eventMap[event_name] = object()
+
+        downlink._finish_downlink_early()
+        assert downlink.satellite.requires_retasking
+        assert not downlink._downlink_empty_event_ready()
+        downlink._clear_action_events()
+
+        downlink.simulator.delete_event.assert_called_once_with(event_name)
+        assert downlink._downlink_empty_event_name is None
+
+    def test_behavior_remains_opt_in(self):
+        assert act.Downlink().variable_duration_downlink is False
+
+
 @pytest.mark.parametrize("kind", ["fsw", "broadcast"])
 @pytest.mark.parametrize(
     "requested,cap,sim_rate,expected_duration",