Fix typing error in SQIL implementation.

HumanCompatibleAI · Jul 18, 2023 · f8584c3 · f8584c3
1 parent 9c5b91c
commit f8584c3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/imitation/algorithms/sqil.py b/src/imitation/algorithms/sqil.py
@@ -181,7 +181,7 @@ def set_demonstrations(
                 next_obs=transition["next_obs"],
                 action=transition["acts"],
                 done=transition["dones"],
-                reward=1,
+                reward=np.array(1.0),
                 infos=[{}],
             )
 
@@ -194,7 +194,7 @@ def add(
         done: np.ndarray,
         infos: List[Dict[str, Any]],
     ) -> None:
-        super().add(obs, next_obs, action, 0, done, infos)
+        super().add(obs, next_obs, action, np.array(0.0), done, infos)
 
     def sample(
         self,