fix some mypy errors

Author: yichao-liang

mypy.ini

@@ -2,7 +2,7 @@
 strict_equality = True
 disallow_untyped_calls = True
 warn_unreachable = True
-exclude = (predicators/envs/assets|venv)
+exclude = (predicators/envs/assets|venv|prompts)
 
 [mypy-predicators.*]
 disallow_untyped_defs = True

predicators/main.py

@@ -233,7 +233,7 @@ def _run_pipeline(env: BaseEnv,
 
 def _handle_offline_learning(
         cogman: CogMan,
-        offline_dataset: Dataset) -> Tuple[float, float, float, dict]:
+        offline_dataset: Dataset) -> Tuple[int, float, float, dict]:
     """Handle offline learning phase and initial evaluation."""
     num_offline_transitions = sum(
         len(traj.actions) for traj in offline_dataset.trajectories)

predicators/utils.py

@@ -4,6 +4,7 @@
 
 import abc
 import contextlib
+import copy
 import datetime
 import functools
 import gc

scripts/plotting/mara/create_latex_table.py

@@ -2,6 +2,7 @@
 
 import math
 from operator import gt, lt
+from typing import Set, Tuple
 
 import numpy as np
 import pandas as pd
@@ -215,7 +216,7 @@ def _main() -> None:
             if abs(mean - best_mean) <= BOLD_NUM_STDS * best_std:
                 bolded.add((outer_label, inner_label))
         # Determine which should be red.
-        red = set()
+        red: Set[Tuple[str, str]] = set()
         # for entry, (mean, std, size) in entry_to_mean_std_size.items():
         #     if size < RED_MIN_SIZE:
         #         red.add(entry)