[rllib] Add type annotations for evaluation/, env/ packages (#9003)

rllib/evaluation/sample_batch_builder.py

@@ -1,17 +1,24 @@
 import collections
 import logging
 import numpy as np
+from typing import List, Any, Dict, Optional, TYPE_CHECKING
 
+from ray.rllib.evaluation.episode import MultiAgentEpisode
+from ray.rllib.policy.policy import Policy
 from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch
 from ray.rllib.utils.annotations import PublicAPI, DeveloperAPI
 from ray.rllib.utils.debug import summarize
+from ray.rllib.utils.types import PolicyID, AgentID
 from ray.rllib.env.base_env import _DUMMY_AGENT_ID
 from ray.util.debug import log_once
 
+if TYPE_CHECKING:
+    from ray.rllib.agents.callbacks import DefaultCallbacks
+
 logger = logging.getLogger(__name__)
 
 
-def to_float_array(v):
+def to_float_array(v: List[Any]) -> np.ndarray:
     arr = np.array(v)
     if arr.dtype == np.float64:
         return arr.astype(np.float32)  # save some memory
@@ -30,11 +37,11 @@ class SampleBatchBuilder:
 
     @PublicAPI
     def __init__(self):
-        self.buffers = collections.defaultdict(list)
+        self.buffers: Dict[str, List] = collections.defaultdict(list)
         self.count = 0
 
     @PublicAPI
-    def add_values(self, **values):
+    def add_values(self, **values: Dict[str, Any]) -> None:
         """Add the given dictionary (row) of values to this batch."""
 
         for k, v in values.items():
@@ -42,7 +49,7 @@ class SampleBatchBuilder:
         self.count += 1
 
     @PublicAPI
-    def add_batch(self, batch):
+    def add_batch(self, batch: SampleBatch) -> None:
         """Add the given batch of values to this batch."""
 
         for k, column in batch.items():
@@ -50,7 +57,7 @@ class SampleBatchBuilder:
         self.count += batch.count
 
     @PublicAPI
-    def build_and_reset(self):
+    def build_and_reset(self) -> SampleBatch:
         """Returns a sample batch including all previously added values."""
 
         batch = SampleBatch(
@@ -75,7 +82,8 @@ class MultiAgentSampleBatchBuilder:
     corresponding policy batch for the agent's policy.
     """
 
-    def __init__(self, policy_map, clip_rewards, callbacks):
+    def __init__(self, policy_map: Dict[PolicyID, Policy], clip_rewards: bool,
+                 callbacks: "DefaultCallbacks"):
         """Initialize a MultiAgentSampleBatchBuilder.
 
         Args:
@@ -102,7 +110,7 @@ class MultiAgentSampleBatchBuilder:
         # Regardless of the number of agents involved in each of these steps.
         self.count = 0
 
-    def total(self):
+    def total(self) -> int:
         """Returns the total number of steps taken in the env (all agents).
 
         Returns:
@@ -112,7 +120,7 @@ class MultiAgentSampleBatchBuilder:
 
         return sum(a.count for a in self.agent_builders.values())
 
-    def has_pending_agent_data(self):
+    def has_pending_agent_data(self) -> bool:
         """Returns whether there is pending unprocessed data.
 
         Returns:
@@ -123,7 +131,8 @@ class MultiAgentSampleBatchBuilder:
         return len(self.agent_builders) > 0
 
     @DeveloperAPI
-    def add_values(self, agent_id, policy_id, **values):
+    def add_values(self, agent_id: AgentID, policy_id: AgentID,
+                   **values: Dict[str, Any]) -> None:
         """Add the given dictionary (row) of values to this batch.
 
         Arguments:
@@ -142,7 +151,8 @@ class MultiAgentSampleBatchBuilder:
 
         self.agent_builders[agent_id].add_values(**values)
 
-    def postprocess_batch_so_far(self, episode=None):
+    def postprocess_batch_so_far(
+            self, episode: Optional[MultiAgentEpisode] = None) -> None:
         """Apply policy postprocessors to any unprocessed rows.
 
         This pushes the postprocessed per-agent batches onto the per-policy
@@ -210,7 +220,7 @@ class MultiAgentSampleBatchBuilder:
         self.agent_builders.clear()
         self.agent_to_policy.clear()
 
-    def check_missing_dones(self):
+    def check_missing_dones(self) -> None:
         for agent_id, builder in self.agent_builders.items():
             if builder.buffers["dones"][-1] is not True:
                 raise ValueError(
@@ -223,7 +233,8 @@ class MultiAgentSampleBatchBuilder:
                     "Alternatively, set no_done_at_end=True to allow this.")
 
     @DeveloperAPI
-    def build_and_reset(self, episode=None):
+    def build_and_reset(self, episode: Optional[MultiAgentEpisode] = None
+                        ) -> MultiAgentBatch:
         """Returns the accumulated sample batches for each policy.
 
         Any unprocessed rows will be first postprocessed with a policy