[rllib] Add type annotations to Trainer class (#8642)

* type trainer

* type it

* fxi

rllib/agents/trainer.py

@@ -6,18 +6,22 @@ import os
 import pickle
 import time
 import tempfile
+from typing import Callable, List, Dict, Union, Any
 
 import ray
 from ray.exceptions import RayError
 from ray.rllib.agents.callbacks import DefaultCallbacks
+from ray.rllib.env import EnvType
 from ray.rllib.env.normalize_actions import NormalizeActionWrapper
 from ray.rllib.models import MODEL_DEFAULTS
+from ray.rllib.policy import Policy, PolicyID
 from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
 from ray.rllib.evaluation.metrics import collect_metrics
 from ray.rllib.optimizers.policy_optimizer import PolicyOptimizer
 from ray.rllib.evaluation.worker_set import WorkerSet
 from ray.rllib.utils import FilterManager, deep_update, merge_dicts
-from ray.rllib.utils.framework import check_framework, try_import_tf
+from ray.rllib.utils.framework import check_framework, try_import_tf, \
+    TensorStructType
 from ray.rllib.utils.annotations import override, PublicAPI, DeveloperAPI
 from ray.rllib.utils.deprecation import DEPRECATED_VALUE, deprecation_warning
 from ray.rllib.utils.from_config import from_config
@@ -25,7 +29,7 @@ from ray.tune.registry import ENV_CREATOR, register_env, _global_registry
 from ray.tune.trainable import Trainable
 from ray.tune.trial import ExportFormat
 from ray.tune.resources import Resources
-from ray.tune.logger import UnifiedLogger
+from ray.tune.logger import Logger, UnifiedLogger
 from ray.tune.result import DEFAULT_RESULTS_DIR
 
 tf = try_import_tf()
@@ -402,7 +406,10 @@ class Trainer(Trainable):
     _override_all_subkeys_if_type_changes = ["exploration_config"]
 
     @PublicAPI
-    def __init__(self, config=None, env=None, logger_creator=None):
+    def __init__(self,
+                 config: dict = None,
+                 env: str = None,
+                 logger_creator: Callable[[], Logger] = None):
         """Initialize an RLLib trainer.
 
         Args:
@@ -446,7 +453,7 @@ class Trainer(Trainable):
 
     @classmethod
     @override(Trainable)
-    def default_resource_request(cls, config):
+    def default_resource_request(cls, config: dict) -> Resources:
         cf = dict(cls._default_config, **config)
         Trainer._validate_config(cf)
         num_workers = cf["num_workers"] + cf["evaluation_num_workers"]
@@ -464,7 +471,7 @@ class Trainer(Trainable):
 
     @override(Trainable)
     @PublicAPI
-    def train(self):
+    def train(self) -> dict:
         """Overrides super.train to synchronize global vars."""
 
         if self._has_policy_optimizer():
@@ -525,14 +532,14 @@ class Trainer(Trainable):
                 workers.local_worker().filters))
 
     @override(Trainable)
-    def _log_result(self, result):
+    def _log_result(self, result: dict):
         self.callbacks.on_train_result(trainer=self, result=result)
         # log after the callback is invoked, so that the user has a chance
         # to mutate the result
         Trainable._log_result(self, result)
 
     @override(Trainable)
-    def _setup(self, config):
+    def _setup(self, config: dict):
         env = self._env_id
         if env:
             config["env"] = env
@@ -649,7 +656,7 @@ class Trainer(Trainable):
             self.optimizer.stop()
 
     @override(Trainable)
-    def _save(self, checkpoint_dir):
+    def _save(self, checkpoint_dir: str) -> str:
         checkpoint_path = os.path.join(checkpoint_dir,
                                        "checkpoint-{}".format(self.iteration))
         pickle.dump(self.__getstate__(), open(checkpoint_path, "wb"))
@@ -657,12 +664,14 @@ class Trainer(Trainable):
         return checkpoint_path
 
     @override(Trainable)
-    def _restore(self, checkpoint_path):
+    def _restore(self, checkpoint_path: str):
         extra_data = pickle.load(open(checkpoint_path, "rb"))
         self.__setstate__(extra_data)
 
     @DeveloperAPI
-    def _make_workers(self, env_creator, policy, config, num_workers):
+    def _make_workers(self, env_creator: Callable[[dict], EnvType],
+                      policy: type, config: dict,
+                      num_workers: int) -> WorkerSet:
         """Default factory method for a WorkerSet running under this Trainer.
 
         Override this method by passing a custom `make_workers` into
@@ -696,7 +705,7 @@ class Trainer(Trainable):
         raise NotImplementedError
 
     @DeveloperAPI
-    def _evaluate(self):
+    def _evaluate(self) -> dict:
         """Evaluates current policy under `evaluation_config` settings.
 
         Note that this default implementation does not do anything beyond
@@ -751,14 +760,14 @@ class Trainer(Trainable):
 
     @PublicAPI
     def compute_action(self,
-                       observation,
-                       state=None,
-                       prev_action=None,
-                       prev_reward=None,
-                       info=None,
-                       policy_id=DEFAULT_POLICY_ID,
-                       full_fetch=False,
-                       explore=None):
+                       observation: TensorStructType,
+                       state: List[Any] = None,
+                       prev_action: TensorStructType = None,
+                       prev_reward: int = None,
+                       info: dict = None,
+                       policy_id: PolicyID = DEFAULT_POLICY_ID,
+                       full_fetch: bool = False,
+                       explore: bool = None) -> TensorStructType:
         """Computes an action for the specified policy on the local Worker.
 
         Note that you can also access the policy object through
@@ -811,17 +820,17 @@ class Trainer(Trainable):
             return result[0]  # backwards compatibility
 
     @property
-    def _name(self):
+    def _name(self) -> str:
         """Subclasses should override this to declare their name."""
         raise NotImplementedError
 
     @property
-    def _default_config(self):
+    def _default_config(self) -> dict:
         """Subclasses should override this to declare their default config."""
         raise NotImplementedError
 
     @PublicAPI
-    def get_policy(self, policy_id=DEFAULT_POLICY_ID):
+    def get_policy(self, policy_id: PolicyID = DEFAULT_POLICY_ID) -> Policy:
         """Return policy for the specified id, or None.
 
         Arguments:
@@ -830,7 +839,7 @@ class Trainer(Trainable):
         return self.workers.local_worker().get_policy(policy_id)
 
     @PublicAPI
-    def get_weights(self, policies=None):
+    def get_weights(self, policies: List[PolicyID] = None) -> dict:
         """Return a dictionary of policy ids to weights.
 
         Arguments:
@@ -840,7 +849,7 @@ class Trainer(Trainable):
         return self.workers.local_worker().get_weights(policies)
 
     @PublicAPI
-    def set_weights(self, weights):
+    def set_weights(self, weights: Dict[PolicyID, dict]):
         """Set policy weights by policy id.
 
         Arguments:
@@ -866,9 +875,9 @@ class Trainer(Trainable):
 
     @DeveloperAPI
     def export_policy_checkpoint(self,
-                                 export_dir,
-                                 filename_prefix="model",
-                                 policy_id=DEFAULT_POLICY_ID):
+                                 export_dir: str,
+                                 filename_prefix: str = "model",
+                                 policy_id: PolicyID = DEFAULT_POLICY_ID):
         """Export tensorflow policy model checkpoint to local directory.
 
         Arguments:
@@ -887,8 +896,8 @@ class Trainer(Trainable):
 
     @DeveloperAPI
     def import_policy_model_from_h5(self,
-                                    import_file,
-                                    policy_id=DEFAULT_POLICY_ID):
+                                    import_file: str,
+                                    policy_id: PolicyID = DEFAULT_POLICY_ID):
         """Imports a policy's model with given policy_id from a local h5 file.
 
         Arguments:
@@ -905,7 +914,8 @@ class Trainer(Trainable):
             import_file, policy_id)
 
     @DeveloperAPI
-    def collect_metrics(self, selected_workers=None):
+    def collect_metrics(self,
+                        selected_workers: List["ActorHandle"] = None) -> dict:
         """Collects metrics from the remote workers of this agent.
 
         This is the same data as returned by a call to train().
@@ -916,14 +926,14 @@ class Trainer(Trainable):
             selected_workers=selected_workers)
 
     @classmethod
-    def resource_help(cls, config):
+    def resource_help(cls, config: dict) -> str:
         return ("\n\nYou can adjust the resource requests of RLlib agents by "
                 "setting `num_workers`, `num_gpus`, and other configs. See "
                 "the DEFAULT_CONFIG defined by each agent for more info.\n\n"
                 "The config of this agent is: {}".format(config))
 
     @classmethod
-    def merge_trainer_configs(cls, config1, config2):
+    def merge_trainer_configs(cls, config1: dict, config2: dict) -> dict:
         config1 = copy.deepcopy(config1)
         # Error if trainer default has deprecated value.
         if config1["sample_batch_size"] != DEPRECATED_VALUE:
@@ -950,7 +960,7 @@ class Trainer(Trainable):
                            cls._override_all_subkeys_if_type_changes)
 
     @staticmethod
-    def _validate_config(config):
+    def _validate_config(config: dict):
         if "policy_graphs" in config["multiagent"]:
             logger.warning(
                 "The `policy_graphs` config has been renamed to `policies`.")
@@ -1037,7 +1047,8 @@ class Trainer(Trainable):
             self.optimizer, PolicyOptimizer)
 
     @override(Trainable)
-    def _export_model(self, export_formats, export_dir):
+    def _export_model(self, export_formats: List[str],
+                      export_dir: str) -> Dict[str, str]:
         ExportFormat.validate(export_formats)
         exported = {}
         if ExportFormat.CHECKPOINT in export_formats:
@@ -1050,7 +1061,7 @@ class Trainer(Trainable):
             exported[ExportFormat.MODEL] = path
         return exported
 
-    def import_model(self, import_file):
+    def import_model(self, import_file: str):
         """Imports a model from import_file.
 
         Note: Currently, only h5 files are supported.
@@ -1075,7 +1086,7 @@ class Trainer(Trainable):
         else:
             return self.import_policy_model_from_h5(import_file)
 
-    def __getstate__(self):
+    def __getstate__(self) -> dict:
         state = {}
         if hasattr(self, "workers"):
             state["worker"] = self.workers.local_worker().save()
@@ -1083,7 +1094,7 @@ class Trainer(Trainable):
             state["optimizer"] = self.optimizer.save()
         return state
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: dict):
         if "worker" in state:
             self.workers.local_worker().restore(state["worker"])
             remote_state = ray.put(state["worker"])
@@ -1092,7 +1103,7 @@ class Trainer(Trainable):
         if "optimizer" in state:
             self.optimizer.restore(state["optimizer"])
 
-    def _register_if_needed(self, env_object):
+    def _register_if_needed(self, env_object: Union[str, EnvType]):
         if isinstance(env_object, str):
             return env_object
         elif isinstance(env_object, type):

rllib/env/__init__.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from ray.rllib.env.base_env import BaseEnv
 from ray.rllib.env.dm_env_wrapper import DMEnv
 from ray.rllib.env.multi_agent_env import MultiAgentEnv
@@ -8,6 +10,9 @@ from ray.rllib.env.env_context import EnvContext
 from ray.rllib.env.policy_client import PolicyClient
 from ray.rllib.env.policy_server_input import PolicyServerInput
 
+# Represents one of the env types in this package.
+EnvType = Any
+
 __all__ = [
     "BaseEnv",
     "MultiAgentEnv",

rllib/utils/framework.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import sys
-from typing import Any
+from typing import Any, Union
 
 from ray.util import log_once
 
@@ -10,6 +10,9 @@ logger = logging.getLogger(__name__)
 # Represents a generic tensor type.
 TensorType = Any
 
+# Either a plain tensor, or a dict or tuple of tensors (or StructTensors).
+TensorStructType = Union[TensorType, dict, tuple]
+
 
 def get_auto_framework():
     """Returns the framework (str) when framework="auto" in the config.