center padding mode, on_interact_step hook

Author: jakegrigsby

amago/experiment.py

@@ -77,7 +77,7 @@ class Experiment:
     :param wandb_group_name: Group runs on wandb dashboard. **Default:** None.
     :param verbose: Print tqdm bars and info to console. **Default:** True.
     :param log_interval: Log extra metrics every N batches. **Default:** 300.
-    :param padded_sampling: Padding for sampling training subsequences. "none", "left", "right", "both". **Default:** "none".
+    :param padded_sampling: Padding for sampling training subsequences. "none", "left", "right", "both", "center". **Default:** "none".
     :param dloader_workers: Number of DataLoader workers for disk loading. Increase for compressed/large trajs.
 
     .. note::
@@ -522,6 +522,15 @@ def policy(self) -> BaseAgent:
         """Returns the current Agent policy free from the accelerator wrapper."""
         return self.accelerator.unwrap_model(self.policy_aclr)
 
+    def on_interact_step(self, obs: dict) -> None:
+        """Hook called before each action during ``self.interact(...)``.
+
+        Override in a subclass to mutate agent state between steps
+        (e.g. resample ``agent._inference_alpha`` at episode boundaries by
+        inspecting ``obs["episode_id"]``).  No-op in the base implementation.
+        """
+        pass
+
     def interact(
         self,
         envs,
@@ -603,6 +612,7 @@ def get_t():
         obs, rl2s, time_idxs = get_t()
         episodes_finished = 0
         for _ in iter_:
+            self.on_interact_step(obs)
             with torch.no_grad():
                 with self.caster():
                     actions, hidden_state = policy.get_actions(

amago/loading.py

@@ -90,6 +90,9 @@ def random_slice(self, length: int, padded_sampling: str = "none"):
                 where the start of the training sequence is not always the first timestep of the trajectory.
                 "right" --> sample while effectively padding the right side of the sequence for cases
                 where the end of the trajectory may be undersampled.
+                "center" --> uniform random window center, clamped to valid range.
+                Always returns a full-length (never padded) window. Fixes endpoint
+                undersampling of "none" when the trajectory is longer than `length`.
         """
         if len(self) <= length:
             start = 0
@@ -101,6 +104,16 @@ def random_slice(self, length: int, padded_sampling: str = "none"):
             start = self._safe_randrange(-length + 1, len(self) - length + 1)
         elif padded_sampling == "right":
             start = self._safe_randrange(0, len(self) - 1)
+        elif padded_sampling == "center":
+            # Uniform center over [0, T), window clamped to fit inside [0, T].
+            # Always length `length`, never padded. Compared to "none", which
+            # severely undersamples the first/last (length - 1) timesteps,
+            # this mode gives them ~length/T marginal coverage instead of
+            # ~1/(T-length+1). Trade-off: interior timesteps still get more
+            # coverage than endpoints when T < 2*length - 1, but the
+            # imbalance is dramatically smaller than "none".
+            center = self._safe_randrange(0, len(self))
+            start = min(max(center - length // 2, 0), len(self) - length)
         else:
             raise ValueError(
                 f"Unrecognized `padded_sampling` mode: `{padded_sampling}`"