Fix mypy errors in trainer code. (#3135)

Author: Chris Elion

.pre-commit-config.yaml

@@ -10,11 +10,7 @@ repos:
  )$
 
 - repo: https://github.com/pre-commit/mirrors-mypy
- rev: v0.750
- # Currently mypy may assert after logging one message. To get all the messages at once, change repo and rev to
- # repo: https://github.com/chriselion/mypy
- # rev: 3d0b6164a9487a6c5cf9d144110b86600fd85e25
- # This is a fork with the assert disabled, although precommit has trouble installing it sometimes.
+ rev: v0.761
  hooks:
  - id: mypy
  name: mypy-ml-agents

ml-agents/mlagents/trainers/ppo/trainer.py

@@ -70,7 +70,7 @@ def __init__(
  self.load = load
  self.multi_gpu = multi_gpu
  self.seed = seed
- self.policy: TFPolicy = None
+ self.policy: PPOPolicy = None # type: ignore
 
  def process_trajectory(self, trajectory: Trajectory) -> None:
  """
@@ -255,6 +255,8 @@ def add_policy(self, name_behavior_id: str, policy: TFPolicy) -> None:
  self.__class__.__name__
  )
  )
+ if not isinstance(policy, PPOPolicy):
+ raise RuntimeError("Non-PPOPolicy passed to PPOTrainer.add_policy()")
  self.policy = policy
 
  def get_policy(self, name_behavior_id: str) -> TFPolicy:

ml-agents/mlagents/trainers/sac/policy.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Mapping
 import numpy as np
 from mlagents.tf_utils import tf
 
@@ -206,7 +206,7 @@ def update(
  return update_stats
 
  def update_reward_signals(
- self, reward_signal_minibatches: Dict[str, Dict], num_sequences: int
+ self, reward_signal_minibatches: Mapping[str, Dict], num_sequences: int
  ) -> Dict[str, float]:
  """
  Only update the reward signals.
@@ -236,7 +236,7 @@ def add_reward_signal_dicts(
  feed_dict: Dict[tf.Tensor, Any],
  update_dict: Dict[str, tf.Tensor],
  stats_needed: Dict[str, str],
- reward_signal_minibatches: Dict[str, Dict],
+ reward_signal_minibatches: Mapping[str, Dict],
  num_sequences: int,
  ) -> None:
  """

ml-agents/mlagents/trainers/sac/trainer.py

@@ -77,7 +77,7 @@ def __init__(
  self.check_param_keys()
  self.load = load
  self.seed = seed
- self.policy: TFPolicy = None
+ self.policy: SACPolicy = None # type: ignore
 
  self.step = 0
  self.train_interval = (
@@ -337,6 +337,8 @@ def add_policy(self, name_behavior_id: str, policy: TFPolicy) -> None:
  self.__class__.__name__
  )
  )
+ if not isinstance(policy, SACPolicy):
+ raise RuntimeError("Non-SACPolicy passed to SACTrainer.add_policy()")
  self.policy = policy
 
  def get_policy(self, name_behavior_id: str) -> TFPolicy: