Merge pull request #109 from davidecoraci/add-dr-estimator

Add DR estimator and RTG regression model

Author: antoine-galataud

hopes/ope/estimators.py

@@ -6,7 +6,7 @@
 
 from hopes.assert_utils import check_array
 from hopes.dev_utils import override
-from hopes.rew.rewards import RegressionBasedRewardModel
+from hopes.rew.rewards import RegressionBasedRewardModel, RTGQModelHGBoost
 
 
 class BaseEstimator(ABC):
@@ -45,8 +45,8 @@ def set_importance_ratios(self, importance_ratios: np.ndarray | None) -> None:
         stickiness correction outside the estimator.
 
         param importance_ratios: Precomputed importance ratios for the logged actions.
-            Supported shapes are ``(n_samples,)`` or ``(n_episodes, steps_per_episode)``.
-        :raises ValueError: If the shape of ``importance_ratios`` is invalid.
+            Supported shapes are `(n_samples,)` or `(n_episodes, steps_per_episode)`.
+        :raises ValueError: If the shape of `importance_ratios` is invalid.
         """
 
         if importance_ratios is None:
@@ -619,17 +619,17 @@ def compute_weighted_rewards(
         step-wise ratios directly instead of recomputing them from policy probabilities.
 
         :param target_policy_action_probabilities: Target policy action probabilities,
-            shape ``(n_samples, n_actions)``.
+            shape `(n_samples, n_actions)`.
         :param behavior_policy_action_probabilities: Behavior policy action probabilities,
-            shape ``(n_samples, n_actions)``.
-        :param rewards: Logged rewards, shape ``(n_samples,)``.
+            shape `(n_samples, n_actions)`.
+        :param rewards: Logged rewards, shape `(n_samples,)`.
         :param steps_per_episode: Number of steps per episode.
-        :param discount_factor: Discount factor in ``[0, 1]``.
+        :param discount_factor: Discount factor in `[0, 1]`.
         :param is_per_decision: Whether to compute per-decision or trajectory-wise weighting.
         :param importance_ratios: Optional precomputed step-wise importance ratios for the
-            logged actions. Supported shapes are ``(n_samples,)`` or
-            ``(n_episodes, steps_per_episode)``.
-        :return: Weighted rewards per episode, shape ``(n_episodes, 1)``.
+            logged actions. Supported shapes are `(n_samples,)` or
+            `(n_episodes, steps_per_episode)`.
+        :return: Weighted rewards per episode, shape `(n_episodes, 1)`.
         """
 
         # rewards, shape: (n, T)
@@ -1067,3 +1067,290 @@ def _bootstrap_sample_policy_value(
         den = float(np.sum(W_b))
 
         return float(num / np.maximum(den, self.eps))
+
+
+class SequentialDoublyRobust(BaseEstimator):
+    r"""Sequential Doubly Robust estimator.
+
+    This estimator computes a per-decision doubly robust estimate using a
+    temporal-difference-style formulation. It combines model-based predictions
+    with cumulative importance weights built from behavior and target policy
+    action probabilities.
+
+    The per-episode estimate is computed as:
+
+    .. math::
+        \hat{V}_{\mathrm{DR}}^{(i)} =
+        \hat{V}(s_{i,0}) +
+        \sum_{t=0}^{T-1}
+        W_{i,t}
+        \left(
+            r_{i,t}
+            + \gamma \hat{V}(s_{i,t+1})
+            - \hat{Q}(s_{i,t}, a_{i,t})
+        \right)
+
+    where:
+
+    .. math::
+        W_{i,t} = \prod_{k=0}^{t} \rho_{i,k}
+
+    and
+
+    .. math::
+        \rho_{i,t} =
+        \frac{\pi_e(a_{i,t} \mid s_{i,t})}{\pi_b(a_{i,t} \mid s_{i,t})}
+
+    with:
+
+    - :math:`i` denoting the episode index,
+    - :math:`t` denoting the timestep index,
+    - :math:`r_{i,t}` the observed reward at timestep :math:`t`,
+    - :math:`\hat{Q}(s_{i,t}, a_{i,t})` the estimated action-value for the logged action,
+    - :math:`\hat{V}(s_{i,t})` the estimated state value under the target policy,
+    - :math:`\gamma` the discount factor,
+    - :math:`\pi_e` the target policy,
+    - :math:`\pi_b` the behavior policy.
+
+    If precomputed step-wise importance ratios are provided, they are used directly.
+    Otherwise, the ratios are constructed from the target and behavior policy action
+    probabilities and the logged actions.
+
+    Stickiness handling, when needed, must be applied upstream during preprocessing.
+    Reference paper: https://arxiv.org/abs/1511.03722
+    """
+
+    def __init__(
+        self,
+        *,
+        steps_per_episode: int,
+        discount_factor: float = 1.0,
+        eps: float = 1e-12,
+        clip: float | None = None,
+    ) -> None:
+        r"""Initialize the Sequential Doubly Robust estimator.
+
+        :param steps_per_episode: Number of timesteps in each episode.
+        :param discount_factor: Discount factor
+        :math:`\gamma` used in the TD correction term. Must be in
+        :math:`[0, 1]`.
+        :param eps: Numerical stabilizer used in importance-ratio computation to avoid
+            division by zero.
+        :param clip: Optional symmetric clipping threshold applied to step-wise importance
+            ratios as
+        :math:`\rho_t \leftarrow \mathrm{clip}(\rho_t, 1 / c, c)`. When provided, it must
+            satisfy
+        :math:`c \geq 1`.
+        """
+        super().__init__()
+
+        assert steps_per_episode > 0, "The number of steps per episode must be positive."
+        assert 0 <= discount_factor <= 1, "The discount factor must be in [0, 1]."
+
+        self.steps_per_episode = steps_per_episode
+        self.discount_factor = discount_factor
+        self.eps = eps
+        self.clip = clip
+
+        self.logged_actions: np.ndarray | None = None
+        self.q_values: np.ndarray | None = None
+
+    def set_logged_actions(self, logged_actions: np.ndarray) -> None:
+        """Set logged actions.
+
+        :param logged_actions: Logged action indices with shape `(n_samples,)`.
+        """
+        self.logged_actions = np.asarray(logged_actions, dtype=np.int64).reshape(-1)
+
+    def set_model_predictions(self, *, q_values: np.ndarray) -> None:
+        r"""Set model-based predictions used by the sequential DR estimator.
+
+        :param q_values: Estimated action-values for all actions, shape `(n_samples,
+            n_actions)`. Each row must contain the estimated action-values
+        :math:`[\hat{Q}(s_t, a)]_{a \in \mathcal{A}}` for the corresponding state.
+        """
+        self.q_values = np.asarray(q_values, dtype=np.float32)
+
+    def fit(
+        self,
+        *,
+        obs_flat: np.ndarray,
+        act_flat: np.ndarray,
+        rew_flat: np.ndarray,
+        num_actions: int,
+        q_model_params: dict | None = None,
+        random_state: int = 0,
+    ) -> RTGQModelHGBoost:
+        """Fit an internal :class:`~hopes.rew.rewards.RTGQModelHGBoost` Q model from raw
+        logged data, then populate :attr:`q_values` and :attr:`logged_actions` automatically.
+
+        This mirrors the design of :meth:`DirectMethod.fit`: rather than building and fitting
+        the Q model externally and injecting the predictions via
+        :meth:`set_model_predictions`, you can pass the raw trajectory data directly and let
+        the estimator handle the model training.  The two workflows remain interchangeable —
+        :meth:`set_model_predictions` is still available for cases where Q-values come from
+        an external source.
+
+        :param obs_flat: Observations, shape `(n_samples, obs_dim)`.
+        :param act_flat: Discrete action indices, shape `(n_samples,)`.
+        :param rew_flat: Rewards, shape `(n_samples,)`.
+        :param num_actions: Total number of discrete actions.
+        :param q_model_params: Optional hyper-parameters forwarded to
+            :class:`~hopes.rew.rewards.RTGQModelHGBoost`.
+        :param random_state: Random seed for the underlying gradient-boosting model.
+        :return: The fitted :class:`~hopes.rew.rewards.RTGQModelHGBoost` instance.
+        """
+        obs_flat = np.asarray(obs_flat, dtype=np.float32)
+        act_flat = np.asarray(act_flat, dtype=np.int64).reshape(-1)
+        rew_flat = np.asarray(rew_flat, dtype=np.float32).reshape(-1)
+
+        q_model = RTGQModelHGBoost(
+            steps_per_episode=self.steps_per_episode,
+            num_actions=num_actions,
+            discount_factor=self.discount_factor,
+            model_params=q_model_params or {},
+            random_state=random_state,
+        )
+        q_model.fit(obs_flat=obs_flat, act_flat=act_flat, rew_flat=rew_flat)
+
+        # Populate q_values for all (state, action) pairs and store logged actions
+        q_values = q_model.predict_q_values(obs_flat=obs_flat)  # shape: (n_samples, num_actions)
+        self.set_model_predictions(q_values=q_values)
+        self.set_logged_actions(act_flat)
+
+        return q_model
+
+    @override(BaseEstimator)
+    def short_name(self) -> str:
+        return "SDR"
+
+    @override(BaseEstimator)
+    def check_parameters(self) -> None:
+        """Check if the estimator parameters are valid."""
+        super().check_parameters()
+
+        if self.eps <= 0:
+            raise ValueError("eps must be > 0.")
+
+        if self.clip is not None and self.clip < 1.0:
+            raise ValueError("clip must be >= 1.0 when provided.")
+
+        n_samples = self.rewards.shape[0]
+        n_actions = self.target_policy_action_probabilities.shape[1]
+
+        if n_samples % self.steps_per_episode != 0:
+            raise ValueError("The number of samples must be divisible by steps_per_episode.")
+
+        if self.q_values is None:
+            raise ValueError("q_values not set. Call set_model_predictions(...).")
+
+        if self.q_values.ndim != 2:
+            raise ValueError("q_values must be a 2D array of shape (n_samples, n_actions).")
+
+        if self.q_values.shape != (n_samples, n_actions):
+            raise ValueError(
+                "q_values must have shape (n_samples, n_actions), matching "
+                "target_policy_action_probabilities."
+            )
+
+        if self.logged_actions is None:
+            raise ValueError("logged_actions must be provided.")
+
+        if self.logged_actions.shape[0] != n_samples:
+            raise ValueError("logged_actions length must match rewards length.")
+
+        if np.any(self.logged_actions < 0) or np.any(self.logged_actions >= n_actions):
+            raise ValueError("logged_actions contains invalid action indices.")
+
+    def _get_stepwise_importance_ratios(self) -> np.ndarray:
+        """Get the step-wise importance ratios.
+
+        :return: Step-wise importance ratios, shape `(n_episodes, steps_per_episode)`.
+        """
+
+        n_samples = self.rewards.shape[0]
+        n_episodes = n_samples // self.steps_per_episode
+
+        if self.importance_ratios is not None:
+            rho = np.asarray(self.importance_ratios, dtype=np.float32)
+            if rho.ndim == 1:
+                rho = rho.reshape(n_episodes, self.steps_per_episode)
+            return rho
+
+        idx = np.arange(n_samples, dtype=np.int64)
+        actions = np.asarray(self.logged_actions, dtype=np.int64).reshape(-1)
+
+        p_e_taken = self.target_policy_action_probabilities[idx, actions].astype(np.float32)
+        p_b_taken = self.behavior_policy_action_probabilities[idx, actions].astype(np.float32)
+
+        rho = p_e_taken / np.maximum(p_b_taken, self.eps)
+
+        if self.clip is not None:
+            rho = np.clip(rho, 1.0 / self.clip, self.clip)
+
+        return rho.reshape(n_episodes, self.steps_per_episode)
+
+    @override(BaseEstimator)
+    def estimate_weighted_rewards(self) -> np.ndarray:
+        """Estimate episode-level sequential DR contributions.
+
+        :return: Episode-level DR estimates, shape `(n_episodes, 1)`.
+        """
+        self.check_parameters()
+
+        # shape: (n_episodes, horizon)
+        rewards = np.asarray(self.rewards, dtype=np.float32).reshape(-1, self.steps_per_episode)
+        # shape: (n_samples, n_actions)
+        q_values = np.asarray(self.q_values, dtype=np.float32)
+        # shape: (n_episodes, horizon)
+        rho = self._get_stepwise_importance_ratios()
+
+        n_episodes, horizon = rewards.shape
+        n_samples = n_episodes * horizon
+
+        # flat index arrays used for advanced indexing into q_values
+        logged_actions = np.asarray(self.logged_actions, dtype=np.int64).reshape(
+            -1
+        )  # shape: (n_samples,)
+        idx = np.arange(n_samples, dtype=np.int64)  # shape: (n_samples,)
+
+        # Q-value of the action actually taken at each (episode, timestep)
+        # shape: (n_episodes, horizon)
+        q_logged = q_values[idx, logged_actions].reshape(n_episodes, horizon)
+
+        # V(s) = Σ_a π_e(a|s) * Q(s, a) — expected value under the target policy
+        # shape: (n_episodes, horizon)
+        v_values = np.sum(
+            self.target_policy_action_probabilities * q_values,
+            axis=1,
+        ).reshape(n_episodes, horizon)
+
+        # V(s_{t+1}): shift v_values one step forward; terminal step gets 0
+        # shape: (n_episodes, horizon)
+        v_next = np.zeros_like(v_values)
+        v_next[:, :-1] = v_values[:, 1:]
+        v_next[:, -1] = 0.0
+
+        # W_{i,t} = Π_{k=0}^{t} ρ_{i,k} — cumulative importance weights up to each timestep
+        # shape: (n_episodes, horizon)
+        cumulative_weights = np.cumprod(rho, axis=1)
+
+        # Initialise each episode's estimate with V(s_0), the model's prediction at the initial state
+        # shape: (n_episodes,)
+        episode_estimates = v_values[:, 0].copy()
+
+        # Accumulate TD-style correction terms across timesteps:
+        #   δ_t = r_t + γ * V(s_{t+1}) - Q(s_t, a_t)   (residual / advantage at step t)
+        #   episode_estimate += W_{i,t} * δ_t            (weighted by cumulative IS ratio)
+        for t in range(horizon):
+            # shape: (n_episodes,)
+            td_correction = rewards[:, t] + self.discount_factor * v_next[:, t] - q_logged[:, t]
+            episode_estimates += cumulative_weights[:, t] * td_correction
+
+        # shape: (n_episodes, 1)
+        return episode_estimates.reshape(-1, 1).astype(np.float32)
+
+    @override(BaseEstimator)
+    def estimate_policy_value(self) -> float:
+        """Estimate the value of the target policy."""
+        return float(np.mean(self.estimate_weighted_rewards()))