Fix type safety and typing issues in AndroidEnv components.

- Annotate log_stream filters to ensure strict typing.
- Refactor specs action_spec construction to use clean, standard loops and assignments, resolving type invariance issues.

PiperOrigin-RevId: 917119228

Author: kenjitoyama

android_env/components/log_stream.py

@@ -26,7 +26,7 @@ class LogStream(metaclass=abc.ABCMeta):
 
   def __init__(self, verbose: bool = False):
     self._verbose = verbose
-    self._filters = []
+    self._filters: list[str] = []
     self._should_stream = threading.Event()
 
   def get_stream_output(self) -> Generator[str, None, None]:

android_env/components/specs.py

@@ -65,36 +65,37 @@ def base_action_spec(
 
   num_actions = len(action_type.ActionType) if enable_key_events else 3
 
-  action_spec = {
-      'action_type':
-          specs.DiscreteArray(num_values=num_actions, name='action_type'),
-      'touch_position':
-          specs.BoundedArray(
-              shape=(2,),
-              dtype=np.float32,
-              minimum=[0.0, 0.0],
-              maximum=[1.0, 1.0],
-              name='touch_position'),
+  action_spec: dict[str, specs.Array] = {
+      'action_type': specs.DiscreteArray(
+          num_values=num_actions, name='action_type'
+      ),
+      'touch_position': specs.BoundedArray(
+          shape=(2,),
+          dtype=np.float32,
+          minimum=[0.0, 0.0],
+          maximum=[1.0, 1.0],
+          name='touch_position',
+      ),
   }
 
+  # Add specs for additional fingers.
   for i in range(2, num_fingers + 1):
-    action_spec.update({
-        f'action_type_{i}':
-            specs.DiscreteArray(
-                num_values=len(action_type.ActionType),
-                name=f'action_type_{i}'),
-        f'touch_position_{i}':
-            specs.BoundedArray(
-                shape=(2,),
-                dtype=np.float32,
-                minimum=[0.0, 0.0],
-                maximum=[1.0, 1.0],
-                name=f'touch_position_{i}'),
-    })
-
+    action_spec[f'action_type_{i}'] = specs.DiscreteArray(
+        num_values=len(action_type.ActionType), name=f'action_type_{i}'
+    )
+    action_spec[f'touch_position_{i}'] = specs.BoundedArray(
+        shape=(2,),
+        dtype=np.float32,
+        minimum=[0.0, 0.0],
+        maximum=[1.0, 1.0],
+        name=f'touch_position_{i}',
+    )
+
+  # Add keycode spec if enabled.
   if enable_key_events:
     action_spec['keycode'] = specs.DiscreteArray(
-        num_values=(1 << 16) - 1, name='keycode')
+        num_values=(1 << 16) - 1, name='keycode'
+    )
 
   return action_spec