cleanup

airboxlab · Apr 19, 2024 · 4f732f6 · 4f732f6
1 parent fb32665
commit 4f732f6
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/doc/source/ope/index.rst b/doc/source/ope/index.rst
@@ -7,9 +7,9 @@ Roadmap
 - [x] Implement Inverse Probability Weighting (IPW) estimator
 - [x] Implement Self-Normalized Inverse Probability Weighting (SNIPW) estimator
 - [x] Implement Direct Method (DM) estimator
-- [ ] Implement Doubly Robust (DR) estimator
 - [X] Implement Trajectory-wise Importance Sampling (TIS) estimator
 - [ ] Implement Per-Decision Importance Sampling (PDIS) estimator
+- [ ] Implement Doubly Robust (DR) estimator
 
 Implemented estimators
 -----------------------

diff --git a/hopes/ope/estimators.py b/hopes/ope/estimators.py
@@ -499,11 +499,11 @@ def estimate_weighted_rewards(self) -> np.ndarray:
         # [gamma^0, gamma^1, ..., gamma^(T-1)] = [gamma^1, gamma^2, ..., gamma^T] / gamma
         discount_factors = np.cumprod(discount_factors, axis=1) / self.discount_factor
 
-        # compute the weighted rewards per trajectory, shape: (n,)
+        # compute the weighted rewards per trajectory, shape: (n, 1)
         weighted_rewards = np.sum(
             importance_weights * rewards * discount_factors,  # (n, 1) * (n, T) * (n, T)
             axis=1,  # sum weights over the trajectory length
-        )
+        ).reshape(-1, 1)
 
         return weighted_rewards
 

diff --git a/tests/test_estimators.py b/tests/test_estimators.py
@@ -164,7 +164,7 @@ def test_tis(self):
 
         wrew = tis.estimate_weighted_rewards()
         self.assertIsInstance(wrew, np.ndarray)
-        self.assertEqual(wrew.shape, (5,))
+        self.assertEqual(wrew.shape, (5, 1))
 
         policy_value = tis.estimate_policy_value()
         self.assertIsInstance(policy_value, float)