From 1e4a1360fd5a5448a3415cfbf6cad2aebb1dc8eb Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Wed, 3 Jun 2020 12:47:35 -0700 Subject: [PATCH] [rllib] Add type annotations to Trainer class (#8642) * type trainer * type it * fxi --- rllib/agents/trainer.py | 87 ++++++++++++++++++++++------------------ rllib/env/__init__.py | 5 +++ rllib/utils/framework.py | 5 ++- 3 files changed, 58 insertions(+), 39 deletions(-) diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 64bb4d93a..90f1d2774 100644 --- a/rllib/agents/trainer.py +++ b/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): diff --git a/rllib/env/__init__.py b/rllib/env/__init__.py index a9483cb7f..f0c87e793 100644 --- a/rllib/env/__init__.py +++ b/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", diff --git a/rllib/utils/framework.py b/rllib/utils/framework.py index c42733789..27770962c 100644 --- a/rllib/utils/framework.py +++ b/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.