Fix pytype and lstm_states indexing

AlignmentResearch · Feb 27, 2024 · b46fce8 · b46fce8
1 parent 09cdebc
commit b46fce8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/stable_baselines3/common/callbacks.py b/stable_baselines3/common/callbacks.py
@@ -440,7 +440,9 @@ def _log_success_callback(self, locals_: Dict[str, Any], globals_: Dict[str, Any
             if maybe_is_success is not None:
                 self._is_success_buffer.append(maybe_is_success)
 
-    def _evaluate_policy(self) -> Union[Tuple[float, float], Tuple[List[float], List[int]]]:
+    def _evaluate_policy(
+        self,
+    ) -> Union[Tuple[float, float], Tuple[List[float], List[int]]]:  # pytype: disable=bad-return-type
         return evaluate_policy(
             self.model,
             self.eval_env,

diff --git a/stable_baselines3/ppo_recurrent/ppo_recurrent.py b/stable_baselines3/ppo_recurrent/ppo_recurrent.py
@@ -213,7 +213,7 @@ def think_for_n_steps(
             return lstm_states
         # ignore because TorchGymObs and TensorTree do not match
         obs_for_start_envs: TorchGymObs = tree_index(obs_tensor, (episode_starts,))  # type: ignore[type-var]
-        lstm_states_for_start_envs = tree_index(lstm_states, (episode_starts,))
+        lstm_states_for_start_envs = tree_index(lstm_states, (slice(None), episode_starts))
         for _ in range(n_steps):
             _, _, _, lstm_states_for_start_envs = self.policy.forward(
                 obs_for_start_envs,