From 1e0e1a45e6a55cbabf8a39b1c045f19e525be1d8 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 19 Jun 2020 13:09:05 -0700 Subject: [PATCH] [rllib] Add type annotations for evaluation/, env/ packages (#9003) --- rllib/agents/callbacks.py | 3 +- rllib/agents/dqn/apex.py | 3 +- rllib/agents/registry.py | 4 +- rllib/agents/trainer.py | 53 ++-- rllib/agents/trainer_template.py | 27 +- rllib/env/__init__.py | 7 +- rllib/env/base_env.py | 80 +++--- rllib/env/env_context.py | 15 +- rllib/env/external_env.py | 35 ++- rllib/env/external_multi_agent_env.py | 29 ++- rllib/env/multi_agent_env.py | 16 +- rllib/env/policy_client.py | 44 ++-- rllib/env/policy_server_input.py | 4 +- rllib/env/remote_vector_env.py | 25 +- rllib/env/unity3d_env.py | 34 ++- rllib/env/vector_env.py | 29 ++- rllib/evaluation/episode.py | 93 ++++--- rllib/evaluation/metrics.py | 30 ++- rllib/evaluation/observation_function.py | 3 +- rllib/evaluation/postprocessing.py | 14 +- rllib/evaluation/rollout_worker.py | 265 +++++++++++-------- rllib/evaluation/sample_batch_builder.py | 35 ++- rllib/evaluation/sampler.py | 318 ++++++++++++++--------- rllib/evaluation/worker_set.py | 51 ++-- rllib/execution/common.py | 9 +- rllib/execution/replay_buffer.py | 2 +- rllib/execution/replay_ops.py | 3 +- rllib/execution/rollout_ops.py | 10 +- rllib/execution/train_ops.py | 4 +- rllib/execution/tree_agg.py | 3 +- rllib/policy/__init__.py | 4 +- rllib/policy/policy.py | 7 - rllib/utils/framework.py | 2 + rllib/utils/types.py | 79 ++++++ 34 files changed, 840 insertions(+), 500 deletions(-) create mode 100644 rllib/utils/types.py diff --git a/rllib/agents/callbacks.py b/rllib/agents/callbacks.py index f4e57789c..7ef06bed7 100644 --- a/rllib/agents/callbacks.py +++ b/rllib/agents/callbacks.py @@ -1,11 +1,12 @@ from typing import Dict from ray.rllib.env import BaseEnv -from ray.rllib.policy import Policy, PolicyID, AgentID +from ray.rllib.policy import Policy from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.evaluation import MultiAgentEpisode, RolloutWorker from ray.rllib.utils.annotations import PublicAPI from ray.rllib.utils.deprecation import deprecation_warning +from ray.rllib.utils.types import AgentID, PolicyID @PublicAPI diff --git a/rllib/agents/dqn/apex.py b/rllib/agents/dqn/apex.py index 463a1cd80..6dfee4766 100644 --- a/rllib/agents/dqn/apex.py +++ b/rllib/agents/dqn/apex.py @@ -6,7 +6,7 @@ from ray.rllib.agents.dqn.dqn import DQNTrainer, \ DEFAULT_CONFIG as DQN_CONFIG, calculate_rr_weights from ray.rllib.agents.dqn.learner_thread import LearnerThread from ray.rllib.execution.common import STEPS_TRAINED_COUNTER, \ - SampleBatchType, _get_shared_metrics, _get_global_vars + _get_shared_metrics, _get_global_vars from ray.rllib.evaluation.worker_set import WorkerSet from ray.rllib.execution.rollout_ops import ParallelRollouts from ray.rllib.execution.concurrency_ops import Concurrently, Enqueue, Dequeue @@ -16,6 +16,7 @@ from ray.rllib.execution.metric_ops import StandardMetricsReporting from ray.rllib.execution.replay_buffer import ReplayActor from ray.rllib.utils import merge_dicts from ray.rllib.utils.actors import create_colocated +from ray.rllib.utils.types import SampleBatchType # yapf: disable # __sphinx_doc_begin__ diff --git a/rllib/agents/registry.py b/rllib/agents/registry.py index 001f348ba..1478237cd 100644 --- a/rllib/agents/registry.py +++ b/rllib/agents/registry.py @@ -117,7 +117,7 @@ ALGORITHMS = { } -def get_agent_class(alg): +def get_agent_class(alg: str) -> type: """Returns the class of a known agent given its name.""" try: @@ -127,7 +127,7 @@ def get_agent_class(alg): return _agent_import_failed(traceback.format_exc()) -def _get_agent_class(alg): +def _get_agent_class(alg: str) -> type: if alg in ALGORITHMS: return ALGORITHMS[alg]() elif alg in CONTRIBUTED_ALGORITHMS: diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 1a4c051a0..6d4774892 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -12,10 +12,10 @@ 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.env.env_context import EnvContext from ray.rllib.models import MODEL_DEFAULTS -from ray.rllib.policy import Policy, PolicyID +from ray.rllib.policy import Policy 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 @@ -26,6 +26,8 @@ from ray.rllib.utils.framework import 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 +from ray.rllib.utils.types import TrainerConfigDict, \ + PartialTrainerConfigDict, EnvInfoDict, ResultDict, EnvType, PolicyID from ray.tune.registry import ENV_CREATOR, register_env, _global_registry from ray.tune.trainable import Trainable from ray.tune.trial import ExportFormat @@ -338,8 +340,9 @@ COMMON_CONFIG = { # === Settings for Multi-Agent Environments === "multiagent": { - # Map from policy ids to tuples of (policy_cls, obs_space, - # act_space, config). See rollout_worker.py for more info. + # Map of type MultiAgentPolicyConfigDict from policy ids to tuples + # of (policy_cls, obs_space, act_space, config). This defines the + # observation and action spaces of the policies and any extra config. "policies": {}, # Function mapping agent ids to policy ids. "policy_mapping_fn": None, @@ -371,13 +374,16 @@ COMMON_CONFIG = { @DeveloperAPI -def with_common_config(extra_config): +def with_common_config( + extra_config: PartialTrainerConfigDict) -> TrainerConfigDict: """Returns the given config dict merged with common agent confs.""" return with_base_config(COMMON_CONFIG, extra_config) -def with_base_config(base_config, extra_config): +def with_base_config( + base_config: TrainerConfigDict, + extra_config: PartialTrainerConfigDict) -> TrainerConfigDict: """Returns the given config dict merged with a base agent conf.""" config = copy.deepcopy(base_config) @@ -418,7 +424,7 @@ class Trainer(Trainable): @PublicAPI def __init__(self, - config: dict = None, + config: TrainerConfigDict = None, env: str = None, logger_creator: Callable[[], Logger] = None): """Initialize an RLLib trainer. @@ -464,7 +470,8 @@ class Trainer(Trainable): @classmethod @override(Trainable) - def default_resource_request(cls, config: dict) -> Resources: + def default_resource_request( + cls, config: PartialTrainerConfigDict) -> Resources: cf = dict(cls._default_config, **config) Trainer._validate_config(cf) num_workers = cf["num_workers"] + cf["evaluation_num_workers"] @@ -482,7 +489,7 @@ class Trainer(Trainable): @override(Trainable) @PublicAPI - def train(self) -> dict: + def train(self) -> ResultDict: """Overrides super.train to synchronize global vars.""" if self._has_policy_optimizer(): @@ -533,7 +540,7 @@ class Trainer(Trainable): return result - def _sync_filters_if_needed(self, workers): + def _sync_filters_if_needed(self, workers: WorkerSet): if self.config.get("observation_filter", "NoFilter") != "NoFilter": FilterManager.synchronize( workers.local_worker().filters, @@ -543,14 +550,14 @@ class Trainer(Trainable): workers.local_worker().filters)) @override(Trainable) - def _log_result(self, result: dict): + def _log_result(self, result: ResultDict): 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: dict): + def _setup(self, config: PartialTrainerConfigDict): env = self._env_id if env: config["env"] = env @@ -678,8 +685,8 @@ class Trainer(Trainable): self.__setstate__(extra_data) @DeveloperAPI - def _make_workers(self, env_creator: Callable[[dict], EnvType], - policy: type, config: dict, + def _make_workers(self, env_creator: Callable[[EnvContext], EnvType], + policy: type, config: TrainerConfigDict, num_workers: int) -> WorkerSet: """Default factory method for a WorkerSet running under this Trainer. @@ -709,7 +716,8 @@ class Trainer(Trainable): logdir=self.logdir) @DeveloperAPI - def _init(self, config, env_creator): + def _init(self, config: TrainerConfigDict, + env_creator: Callable[[EnvContext], EnvType]): """Subclasses should override this for custom initialization.""" raise NotImplementedError @@ -773,7 +781,7 @@ class Trainer(Trainable): state: List[Any] = None, prev_action: TensorStructType = None, prev_reward: int = None, - info: dict = None, + info: EnvInfoDict = None, policy_id: PolicyID = DEFAULT_POLICY_ID, full_fetch: bool = False, explore: bool = None) -> TensorStructType: @@ -923,7 +931,7 @@ class Trainer(Trainable): raise NotImplementedError @property - def _default_config(self) -> dict: + def _default_config(self) -> TrainerConfigDict: """Subclasses should override this to declare their default config.""" raise NotImplementedError @@ -956,7 +964,9 @@ class Trainer(Trainable): self.workers.local_worker().set_weights(weights) @DeveloperAPI - def export_policy_model(self, export_dir, policy_id=DEFAULT_POLICY_ID): + def export_policy_model(self, + export_dir: str, + policy_id: PolicyID = DEFAULT_POLICY_ID): """Export policy model with given policy_id to local directory. Arguments: @@ -1024,14 +1034,15 @@ class Trainer(Trainable): selected_workers=selected_workers) @classmethod - def resource_help(cls, config: dict) -> str: + def resource_help(cls, config: TrainerConfigDict) -> 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: dict, config2: dict) -> dict: + def merge_trainer_configs(cls, config1: TrainerConfigDict, + config2: PartialTrainerConfigDict) -> dict: config1 = copy.deepcopy(config1) # Error if trainer default has deprecated value. if config1["sample_batch_size"] != DEPRECATED_VALUE: @@ -1058,7 +1069,7 @@ class Trainer(Trainable): cls._override_all_subkeys_if_type_changes) @staticmethod - def _validate_config(config: dict): + def _validate_config(config: PartialTrainerConfigDict): if "policy_graphs" in config["multiagent"]: deprecation_warning("policy_graphs", "policies") # Backwards compatibility. diff --git a/rllib/agents/trainer_template.py b/rllib/agents/trainer_template.py index c99e027e0..e94110d70 100644 --- a/rllib/agents/trainer_template.py +++ b/rllib/agents/trainer_template.py @@ -1,18 +1,22 @@ +from typing import Callable, Optional, List, Iterable import logging import time from ray.rllib.agents.trainer import Trainer, COMMON_CONFIG +from ray.rllib.evaluation.worker_set import WorkerSet from ray.rllib.execution.rollout_ops import ParallelRollouts, ConcatBatches from ray.rllib.execution.train_ops import TrainOneStep from ray.rllib.execution.metric_ops import StandardMetricsReporting +from ray.rllib.policy import Policy from ray.rllib.utils import add_mixins from ray.rllib.utils.annotations import override, DeveloperAPI from ray.rllib.utils.deprecation import deprecation_warning +from ray.rllib.utils.types import TrainerConfigDict, ResultDict logger = logging.getLogger(__name__) -def default_execution_plan(workers, config): +def default_execution_plan(workers: WorkerSet, config: TrainerConfigDict): # Collects experiences in parallel from multiple RolloutWorker actors. rollouts = ParallelRollouts(workers, mode="bulk_sync") @@ -30,23 +34,24 @@ def default_execution_plan(workers, config): @DeveloperAPI def build_trainer( - name, - default_policy, - default_config=None, - validate_config=None, + name: str, + default_policy: Optional[Policy], + default_config: TrainerConfigDict = None, + validate_config: Callable[[TrainerConfigDict], None] = None, get_initial_state=None, # DEPRECATED - get_policy_class=None, - before_init=None, + get_policy_class: Callable[[TrainerConfigDict], Policy] = None, + before_init: Callable[[Trainer], None] = None, make_workers=None, # DEPRECATED make_policy_optimizer=None, # DEPRECATED - after_init=None, + after_init: Callable[[Trainer], None] = None, before_train_step=None, # DEPRECATED after_optimizer_step=None, # DEPRECATED after_train_result=None, # DEPRECATED collect_metrics_fn=None, # DEPRECATED - before_evaluate_fn=None, - mixins=None, - execution_plan=default_execution_plan): + before_evaluate_fn: Callable[[Trainer], None] = None, + mixins: List[type] = None, + execution_plan: Callable[[WorkerSet, TrainerConfigDict], Iterable[ + ResultDict]] = default_execution_plan): """Helper function for defining a custom trainer. Functions will be run in this order to initialize the trainer: diff --git a/rllib/env/__init__.py b/rllib/env/__init__.py index f0c87e793..5976b21c6 100644 --- a/rllib/env/__init__.py +++ b/rllib/env/__init__.py @@ -1,7 +1,6 @@ -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.unity3d_env import Unity3DEnv from ray.rllib.env.multi_agent_env import MultiAgentEnv from ray.rllib.env.external_env import ExternalEnv from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv @@ -10,9 +9,6 @@ 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", @@ -21,6 +17,7 @@ __all__ = [ "VectorEnv", "EnvContext", "DMEnv", + "Unity3DEnv", "PolicyClient", "PolicyServerInput", ] diff --git a/rllib/env/base_env.py b/rllib/env/base_env.py index 4c41986c9..43389bff4 100644 --- a/rllib/env/base_env.py +++ b/rllib/env/base_env.py @@ -1,8 +1,15 @@ +from typing import Callable, Tuple, Optional, List, Dict, Any, TYPE_CHECKING + from ray.rllib.env.external_env import ExternalEnv from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv from ray.rllib.env.multi_agent_env import MultiAgentEnv from ray.rllib.env.vector_env import VectorEnv from ray.rllib.utils.annotations import override, PublicAPI +from ray.rllib.utils.types import EnvType, MultiEnvDict, EnvID, \ + AgentID, MultiAgentDict + +if TYPE_CHECKING: + from ray.rllib.models.preprocessors import Preprocessor ASYNC_RESET_RETURN = "async_reset_return" @@ -73,11 +80,11 @@ class BaseEnv: """ @staticmethod - def to_base_env(env, - make_env=None, - num_envs=1, - remote_envs=False, - remote_env_batch_wait_ms=0): + def to_base_env(env: EnvType, + make_env: Callable[[int], EnvType] = None, + num_envs: int = 1, + remote_envs: bool = False, + remote_env_batch_wait_ms: bool = 0) -> "BaseEnv": """Wraps any env type as needed to expose the async interface.""" from ray.rllib.env.remote_vector_env import RemoteVectorEnv @@ -128,7 +135,8 @@ class BaseEnv: return env @PublicAPI - def poll(self): + def poll(self) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, + MultiEnvDict, MultiEnvDict]: """Returns observations from ready agents. The returns are two-level dicts mapping from env_id to a dict of @@ -151,7 +159,7 @@ class BaseEnv: raise NotImplementedError @PublicAPI - def send_actions(self, action_dict): + def send_actions(self, action_dict: MultiEnvDict) -> None: """Called to send actions back to running agents in this env. Actions should be sent for each ready agent that returned observations @@ -163,7 +171,8 @@ class BaseEnv: raise NotImplementedError @PublicAPI - def try_reset(self, env_id=None): + def try_reset(self, + env_id: Optional[EnvID] = None) -> Optional[MultiAgentDict]: """Attempt to reset the sub-env with the given id or all sub-envs. If the environment does not support synchronous reset, None can be @@ -179,7 +188,7 @@ class BaseEnv: return None @PublicAPI - def get_unwrapped(self): + def get_unwrapped(self) -> List[EnvType]: """Return a reference to the underlying gym envs, if any. Returns: @@ -188,7 +197,7 @@ class BaseEnv: return [] @PublicAPI - def stop(self): + def stop(self) -> None: """Releases all resources used.""" for env in self.get_unwrapped(): @@ -200,14 +209,18 @@ class BaseEnv: _DUMMY_AGENT_ID = "agent0" -def _with_dummy_agent_id(env_id_to_values, dummy_id=_DUMMY_AGENT_ID): +def _with_dummy_agent_id(env_id_to_values: Dict[EnvID, Any], + dummy_id: "AgentID" = _DUMMY_AGENT_ID + ) -> MultiEnvDict: return {k: {dummy_id: v} for (k, v) in env_id_to_values.items()} class _ExternalEnvToBaseEnv(BaseEnv): """Internal adapter of ExternalEnv to BaseEnv.""" - def __init__(self, external_env, preprocessor=None): + def __init__(self, + external_env: ExternalEnv, + preprocessor: "Preprocessor" = None): self.external_env = external_env self.prep = preprocessor self.multiagent = issubclass(type(external_env), ExternalMultiAgentEnv) @@ -219,7 +232,8 @@ class _ExternalEnvToBaseEnv(BaseEnv): external_env.start() @override(BaseEnv) - def poll(self): + def poll(self) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, + MultiEnvDict, MultiEnvDict]: with self.external_env._results_avail_condition: results = self._poll() while len(results[0]) == 0: @@ -234,7 +248,7 @@ class _ExternalEnvToBaseEnv(BaseEnv): return results @override(BaseEnv) - def send_actions(self, action_dict): + def send_actions(self, action_dict: MultiEnvDict) -> None: if self.multiagent: for env_id, actions in action_dict.items(): self.external_env._episodes[env_id].action_queue.put(actions) @@ -243,7 +257,8 @@ class _ExternalEnvToBaseEnv(BaseEnv): self.external_env._episodes[env_id].action_queue.put( action[_DUMMY_AGENT_ID]) - def _poll(self): + def _poll(self) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, + MultiEnvDict, MultiEnvDict]: all_obs, all_rewards, all_dones, all_infos = {}, {}, {}, {} off_policy_actions = {} for eid, episode in self.external_env._episodes.copy().items(): @@ -293,7 +308,7 @@ class _VectorEnvToBaseEnv(BaseEnv): environments before calling send_actions(). """ - def __init__(self, vector_env): + def __init__(self, vector_env: VectorEnv): self.vector_env = vector_env self.action_space = vector_env.action_space self.observation_space = vector_env.observation_space @@ -304,7 +319,8 @@ class _VectorEnvToBaseEnv(BaseEnv): self.cur_infos = [None for _ in range(self.num_envs)] @override(BaseEnv) - def poll(self): + def poll(self) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, + MultiEnvDict, MultiEnvDict]: if self.new_obs is None: self.new_obs = self.vector_env.vector_reset() new_obs = dict(enumerate(self.new_obs)) @@ -321,7 +337,7 @@ class _VectorEnvToBaseEnv(BaseEnv): _with_dummy_agent_id(infos), {} @override(BaseEnv) - def send_actions(self, action_dict): + def send_actions(self, action_dict: MultiEnvDict) -> None: action_vector = [None] * self.num_envs for i in range(self.num_envs): action_vector[i] = action_dict[i][_DUMMY_AGENT_ID] @@ -329,11 +345,12 @@ class _VectorEnvToBaseEnv(BaseEnv): self.vector_env.vector_step(action_vector) @override(BaseEnv) - def try_reset(self, env_id): + def try_reset(self, + env_id: Optional[EnvID] = None) -> Optional[MultiAgentDict]: return {_DUMMY_AGENT_ID: self.vector_env.reset_at(env_id)} @override(BaseEnv) - def get_unwrapped(self): + def get_unwrapped(self) -> List[EnvType]: return self.vector_env.get_unwrapped() @@ -343,7 +360,8 @@ class _MultiAgentEnvToBaseEnv(BaseEnv): This also supports vectorization if num_envs > 1. """ - def __init__(self, make_env, existing_envs, num_envs): + def __init__(self, make_env: Callable[[int], EnvType], + existing_envs: List[MultiAgentEnv], num_envs: int): """Wrap existing multi-agent envs. Arguments: @@ -364,14 +382,15 @@ class _MultiAgentEnvToBaseEnv(BaseEnv): self.env_states = [_MultiAgentEnvState(env) for env in self.envs] @override(BaseEnv) - def poll(self): + def poll(self) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, + MultiEnvDict, MultiEnvDict]: obs, rewards, dones, infos = {}, {}, {}, {} for i, env_state in enumerate(self.env_states): obs[i], rewards[i], dones[i], infos[i] = env_state.poll() return obs, rewards, dones, infos, {} @override(BaseEnv) - def send_actions(self, action_dict): + def send_actions(self, action_dict: MultiEnvDict) -> None: for env_id, agent_dict in action_dict.items(): if env_id in self.dones: raise ValueError("Env {} is already done".format(env_id)) @@ -397,7 +416,8 @@ class _MultiAgentEnvToBaseEnv(BaseEnv): self.env_states[env_id].observe(obs, rewards, dones, infos) @override(BaseEnv) - def try_reset(self, env_id): + def try_reset(self, + env_id: Optional[EnvID] = None) -> Optional[MultiAgentDict]: obs = self.env_states[env_id].reset() assert isinstance(obs, dict), "Not a multi-agent obs" if obs is not None and env_id in self.dones: @@ -405,17 +425,18 @@ class _MultiAgentEnvToBaseEnv(BaseEnv): return obs @override(BaseEnv) - def get_unwrapped(self): + def get_unwrapped(self) -> List[EnvType]: return [state.env for state in self.env_states] class _MultiAgentEnvState: - def __init__(self, env): + def __init__(self, env: MultiAgentEnv): assert isinstance(env, MultiAgentEnv) self.env = env self.initialized = False - def poll(self): + def poll(self) -> Tuple[MultiAgentDict, MultiAgentDict, MultiAgentDict, + MultiAgentDict, MultiAgentDict]: if not self.initialized: self.reset() self.initialized = True @@ -427,13 +448,14 @@ class _MultiAgentEnvState: self.last_infos = {} return obs, rew, dones, info - def observe(self, obs, rewards, dones, infos): + def observe(self, obs: MultiAgentDict, rewards: MultiAgentDict, + dones: MultiAgentDict, infos: MultiAgentDict): self.last_obs = obs self.last_rewards = rewards self.last_dones = dones self.last_infos = infos - def reset(self): + def reset(self) -> MultiAgentDict: self.last_obs = self.env.reset() self.last_rewards = { agent_id: None diff --git a/rllib/env/env_context.py b/rllib/env/env_context.py index 1507cd51f..e8ed98b93 100644 --- a/rllib/env/env_context.py +++ b/rllib/env/env_context.py @@ -1,4 +1,5 @@ from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.types import EnvConfigDict @PublicAPI @@ -19,17 +20,21 @@ class EnvContext(dict): remote (bool): Whether environment should be remote or not. """ - def __init__(self, env_config, worker_index, vector_index=0, remote=False): + def __init__(self, + env_config: EnvConfigDict, + worker_index: int, + vector_index: int = 0, + remote: bool = False): dict.__init__(self, env_config) self.worker_index = worker_index self.vector_index = vector_index self.remote = remote def copy_with_overrides(self, - env_config=None, - worker_index=None, - vector_index=None, - remote=None): + env_config: EnvConfigDict = None, + worker_index: int = None, + vector_index: int = None, + remote: bool = None): return EnvContext( env_config if env_config is not None else self, worker_index if worker_index is not None else self.worker_index, diff --git a/rllib/env/external_env.py b/rllib/env/external_env.py index 092b32978..fdac97fb1 100644 --- a/rllib/env/external_env.py +++ b/rllib/env/external_env.py @@ -1,8 +1,11 @@ from six.moves import queue +import gym import threading import uuid +from typing import Optional from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.types import EnvActionType, EnvObsType, EnvInfoDict @PublicAPI @@ -36,7 +39,10 @@ class ExternalEnv(threading.Thread): """ @PublicAPI - def __init__(self, action_space, observation_space, max_concurrent=100): + def __init__(self, + action_space: gym.Space, + observation_space: gym.Space, + max_concurrent: int = 100): """Initializes an external env. Args: @@ -74,7 +80,9 @@ class ExternalEnv(threading.Thread): raise NotImplementedError @PublicAPI - def start_episode(self, episode_id=None, training_enabled=True): + def start_episode(self, + episode_id: Optional[str] = None, + training_enabled: bool = True) -> str: """Record the start of an episode. Args: @@ -104,7 +112,8 @@ class ExternalEnv(threading.Thread): return episode_id @PublicAPI - def get_action(self, episode_id, observation): + def get_action(self, episode_id: str, + observation: EnvObsType) -> EnvActionType: """Record an observation and get the on-policy action. Args: @@ -119,7 +128,8 @@ class ExternalEnv(threading.Thread): return episode.wait_for_action(observation) @PublicAPI - def log_action(self, episode_id, observation, action): + def log_action(self, episode_id: str, observation: EnvObsType, + action: EnvActionType) -> None: """Record an observation and (off-policy) action taken. Args: @@ -132,7 +142,10 @@ class ExternalEnv(threading.Thread): episode.log_action(observation, action) @PublicAPI - def log_returns(self, episode_id, reward, info=None): + def log_returns(self, + episode_id: str, + reward: float, + info: EnvInfoDict = None) -> None: """Record returns from the environment. The reward will be attributed to the previous action taken by the @@ -152,7 +165,7 @@ class ExternalEnv(threading.Thread): episode.cur_info = info or {} @PublicAPI - def end_episode(self, episode_id, observation): + def end_episode(self, episode_id: str, observation: EnvObsType) -> None: """Record the end of an episode. Args: @@ -164,7 +177,7 @@ class ExternalEnv(threading.Thread): self._finished.add(episode.episode_id) episode.done(observation) - def _get(self, episode_id): + def _get(self, episode_id: str) -> "_ExternalEnvEpisode": """Get a started episode or raise an error.""" if episode_id in self._finished: @@ -181,10 +194,10 @@ class _ExternalEnvEpisode: """Tracked state for each active episode.""" def __init__(self, - episode_id, - results_avail_condition, - training_enabled, - multiagent=False): + episode_id: str, + results_avail_condition: threading.Condition, + training_enabled: bool, + multiagent: bool = False): self.episode_id = episode_id self.results_avail_condition = results_avail_condition self.training_enabled = training_enabled diff --git a/rllib/env/external_multi_agent_env.py b/rllib/env/external_multi_agent_env.py index a7d930461..775888341 100644 --- a/rllib/env/external_multi_agent_env.py +++ b/rllib/env/external_multi_agent_env.py @@ -1,7 +1,10 @@ import uuid +import gym +from typing import Optional from ray.rllib.utils.annotations import override, PublicAPI from ray.rllib.env.external_env import ExternalEnv, _ExternalEnvEpisode +from ray.rllib.utils.types import MultiAgentDict @PublicAPI @@ -9,7 +12,10 @@ class ExternalMultiAgentEnv(ExternalEnv): """This is the multi-agent version of ExternalEnv.""" @PublicAPI - def __init__(self, action_space, observation_space, max_concurrent=100): + def __init__(self, + action_space: gym.Space, + observation_space: gym.Space, + max_concurrent: int = 100): """Initialize a multi-agent external env. ExternalMultiAgentEnv subclasses must call this during their __init__. @@ -51,7 +57,9 @@ class ExternalMultiAgentEnv(ExternalEnv): @PublicAPI @override(ExternalEnv) - def start_episode(self, episode_id=None, training_enabled=True): + def start_episode(self, + episode_id: Optional[str] = None, + training_enabled: bool = True) -> str: if episode_id is None: episode_id = uuid.uuid4().hex @@ -73,7 +81,8 @@ class ExternalMultiAgentEnv(ExternalEnv): @PublicAPI @override(ExternalEnv) - def get_action(self, episode_id, observation_dict): + def get_action(self, episode_id: str, + observation_dict: MultiAgentDict) -> MultiAgentDict: """Record an observation and get the on-policy action. observation_dict is expected to contain the observation of all agents acting in this episode step. @@ -91,7 +100,8 @@ class ExternalMultiAgentEnv(ExternalEnv): @PublicAPI @override(ExternalEnv) - def log_action(self, episode_id, observation_dict, action_dict): + def log_action(self, episode_id: str, observation_dict: MultiAgentDict, + action_dict: MultiAgentDict) -> None: """Record an observation and (off-policy) action taken. Arguments: @@ -106,10 +116,10 @@ class ExternalMultiAgentEnv(ExternalEnv): @PublicAPI @override(ExternalEnv) def log_returns(self, - episode_id, - reward_dict, - info_dict=None, - multiagent_done_dict=None): + episode_id: str, + reward_dict: MultiAgentDict, + info_dict: MultiAgentDict = None, + multiagent_done_dict: MultiAgentDict = None) -> None: """Record returns from the environment. The reward will be attributed to the previous action taken by the @@ -142,7 +152,8 @@ class ExternalMultiAgentEnv(ExternalEnv): @PublicAPI @override(ExternalEnv) - def end_episode(self, episode_id, observation_dict): + def end_episode(self, episode_id: str, + observation_dict: MultiAgentDict) -> None: """Record the end of an episode. Arguments: diff --git a/rllib/env/multi_agent_env.py b/rllib/env/multi_agent_env.py index 214ea754d..4034fb81c 100644 --- a/rllib/env/multi_agent_env.py +++ b/rllib/env/multi_agent_env.py @@ -1,4 +1,8 @@ +from typing import Tuple, Dict, List +import gym + from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.types import MultiAgentDict, AgentID # If the obs space is Dict type, look for the global state under this key. ENV_STATE = "state" @@ -44,7 +48,7 @@ class MultiAgentEnv: """ @PublicAPI - def reset(self): + def reset(self) -> MultiAgentDict: """Resets the env and returns observations from ready agents. Returns: @@ -53,7 +57,9 @@ class MultiAgentEnv: raise NotImplementedError @PublicAPI - def step(self, action_dict): + def step( + self, action_dict: MultiAgentDict + ) -> Tuple[MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict]: """Returns observations from ready agents. The returns are dicts mapping from agent_id strings to values. The @@ -73,7 +79,11 @@ class MultiAgentEnv: # yapf: disable # __grouping_doc_begin__ @PublicAPI - def with_agent_groups(self, groups, obs_space=None, act_space=None): + def with_agent_groups( + self, + groups: Dict[str, List[AgentID]], + obs_space: gym.Space = None, + act_space: gym.Space = None) -> "MultiAgentEnv": """Convenience method for grouping together agents in this env. An agent group is a list of agent ids that are mapped to a single diff --git a/rllib/env/policy_client.py b/rllib/env/policy_client.py index 958775360..67835efde 100644 --- a/rllib/env/policy_client.py +++ b/rllib/env/policy_client.py @@ -7,12 +7,15 @@ inference is faster but causes more compute to be done on the client. import logging import threading import time +from typing import Union, Optional import ray.cloudpickle as pickle from ray.rllib.evaluation.rollout_worker import RolloutWorker from ray.rllib.env import ExternalEnv, MultiAgentEnv, ExternalMultiAgentEnv from ray.rllib.policy.sample_batch import MultiAgentBatch from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.types import MultiAgentDict, EnvInfoDict, EnvObsType, \ + EnvActionType logger = logging.getLogger(__name__) logger.setLevel("INFO") # TODO(ekl) seems to be needed for cartpole_client.py @@ -43,7 +46,10 @@ class PolicyClient: END_EPISODE = "END_EPISODE" @PublicAPI - def __init__(self, address, inference_mode="local", update_interval=10.0): + def __init__(self, + address: str, + inference_mode: str = "local", + update_interval: float = 10.0): """Create a PolicyClient instance. Args: @@ -66,7 +72,9 @@ class PolicyClient: "inference_mode must be either 'local' or 'remote'") @PublicAPI - def start_episode(self, episode_id=None, training_enabled=True): + def start_episode(self, + episode_id: Optional[str] = None, + training_enabled: bool = True) -> str: """Record the start of one or more episode(s). Args: @@ -90,7 +98,9 @@ class PolicyClient: })["episode_id"] @PublicAPI - def get_action(self, episode_id, observation): + def get_action(self, episode_id: str, + observation: Union[EnvObsType, MultiAgentDict] + ) -> Union[EnvActionType, MultiAgentDict]: """Record an observation and get the on-policy action. Arguments: @@ -119,7 +129,9 @@ class PolicyClient: })["action"] @PublicAPI - def log_action(self, episode_id, observation, action): + def log_action(self, episode_id: str, + observation: Union[EnvObsType, MultiAgentDict], + action: Union[EnvActionType, MultiAgentDict]) -> None: """Record an observation and (off-policy) action taken. Arguments: @@ -140,11 +152,12 @@ class PolicyClient: }) @PublicAPI - def log_returns(self, - episode_id, - reward, - info=None, - multiagent_done_dict=None): + def log_returns( + self, + episode_id: str, + reward: int, + info: Union[EnvInfoDict, MultiAgentDict] = None, + multiagent_done_dict: Optional[MultiAgentDict] = None) -> None: """Record returns from the environment. The reward will be attributed to the previous action taken by the @@ -175,7 +188,8 @@ class PolicyClient: }) @PublicAPI - def end_episode(self, episode_id, observation): + def end_episode(self, episode_id: str, + observation: Union[EnvObsType, MultiAgentDict]) -> None: """Record the end of an episode. Arguments: @@ -194,7 +208,7 @@ class PolicyClient: }) @PublicAPI - def update_policy_weights(self): + def update_policy_weights(self) -> None: """Query the server for new policy weights, if local inference is enabled. """ self._update_local_policy(force=True) @@ -217,7 +231,7 @@ class PolicyClient: "command": PolicyClient.GET_WORKER_ARGS, })["worker_args"] (self.rollout_worker, - self.inference_thread) = create_embedded_rollout_worker( + self.inference_thread) = _create_embedded_rollout_worker( kwargs, self._send) self.env = self.rollout_worker.env @@ -270,7 +284,7 @@ class _LocalInferenceThread(threading.Thread): logger.info("Error: inference worker thread died!", e) -def auto_wrap_external(real_env_creator): +def _auto_wrap_external(real_env_creator): """Wrap an environment in the ExternalEnv interface if needed. Args: @@ -307,7 +321,7 @@ def auto_wrap_external(real_env_creator): return wrapped_creator -def create_embedded_rollout_worker(kwargs, send_fn): +def _create_embedded_rollout_worker(kwargs, send_fn): """Create a local rollout worker and a thread that samples from it. Arguments: @@ -321,7 +335,7 @@ def create_embedded_rollout_worker(kwargs, send_fn): del kwargs["input_creator"] logger.info("Creating rollout worker with kwargs={}".format(kwargs)) real_env_creator = kwargs["env_creator"] - kwargs["env_creator"] = auto_wrap_external(real_env_creator) + kwargs["env_creator"] = _auto_wrap_external(real_env_creator) rollout_worker = RolloutWorker(**kwargs) inference_thread = _LocalInferenceThread(rollout_worker, send_fn) diff --git a/rllib/env/policy_server_input.py b/rllib/env/policy_server_input.py index 2d8a4770d..60df49d2a 100644 --- a/rllib/env/policy_server_input.py +++ b/rllib/env/policy_server_input.py @@ -9,7 +9,7 @@ from socketserver import ThreadingMixIn import ray.cloudpickle as pickle from ray.rllib.offline.input_reader import InputReader from ray.rllib.env.policy_client import PolicyClient, \ - create_embedded_rollout_worker + _create_embedded_rollout_worker from ray.rllib.utils.annotations import override, PublicAPI logger = logging.getLogger(__name__) @@ -114,7 +114,7 @@ def _make_handler(rollout_worker, samples_queue, metrics_queue): with lock: if child_rollout_worker is None: (child_rollout_worker, - inference_thread) = create_embedded_rollout_worker( + inference_thread) = _create_embedded_rollout_worker( rollout_worker.creation_args(), report_data) child_rollout_worker.set_weights(rollout_worker.get_weights()) diff --git a/rllib/env/remote_vector_env.py b/rllib/env/remote_vector_env.py index f56653c39..8456a905e 100644 --- a/rllib/env/remote_vector_env.py +++ b/rllib/env/remote_vector_env.py @@ -1,21 +1,28 @@ import logging +from typing import Tuple, Callable, Optional import ray from ray.rllib.env.base_env import BaseEnv, _DUMMY_AGENT_ID, ASYNC_RESET_RETURN +from ray.rllib.utils.annotations import override, PublicAPI +from ray.rllib.utils.types import MultiEnvDict, EnvType, EnvID, MultiAgentDict logger = logging.getLogger(__name__) +@PublicAPI class RemoteVectorEnv(BaseEnv): """Vector env that executes envs in remote workers. This provides dynamic batching of inference as observations are returned from the remote simulator actors. Both single and multi-agent child envs are supported, and envs can be stepped synchronously or async. + + You shouldn't need to instantiate this class directly. It's automatically + inserted when you use the `remote_worker_envs` option for Trainers. """ - def __init__(self, make_env, num_envs, multiagent, - remote_env_batch_wait_ms): + def __init__(self, make_env: Callable[[int], EnvType], num_envs: int, + multiagent: bool, remote_env_batch_wait_ms: int): self.make_local_env = make_env self.num_envs = num_envs self.multiagent = multiagent @@ -24,7 +31,9 @@ class RemoteVectorEnv(BaseEnv): self.actors = None # lazy init self.pending = None # lazy init - def poll(self): + @override(BaseEnv) + def poll(self) -> Tuple[MultiEnvDict, MultiEnvDict, MultiEnvDict, + MultiEnvDict, MultiEnvDict]: if self.actors is None: def make_remote_env(i): @@ -65,19 +74,23 @@ class RemoteVectorEnv(BaseEnv): logger.debug("Got obs batch for actors {}".format(env_ids)) return obs, rewards, dones, infos, {} - def send_actions(self, action_dict): + @PublicAPI + def send_actions(self, action_dict: MultiEnvDict) -> None: for env_id, actions in action_dict.items(): actor = self.actors[env_id] obj_id = actor.step.remote(actions) self.pending[obj_id] = actor - def try_reset(self, env_id): + @PublicAPI + def try_reset(self, + env_id: Optional[EnvID] = None) -> Optional[MultiAgentDict]: actor = self.actors[env_id] obj_id = actor.reset.remote() self.pending[obj_id] = actor return ASYNC_RESET_RETURN - def stop(self): + @PublicAPI + def stop(self) -> None: if self.actors is not None: for actor in self.actors: actor.__ray_terminate__.remote() diff --git a/rllib/env/unity3d_env.py b/rllib/env/unity3d_env.py index 66c7337d0..809b94de9 100644 --- a/rllib/env/unity3d_env.py +++ b/rllib/env/unity3d_env.py @@ -1,11 +1,11 @@ -from gym.spaces import Box, MultiDiscrete, Tuple +from gym.spaces import Box, MultiDiscrete, Tuple as TupleSpace import logging -import mlagents_envs -from mlagents_envs.environment import UnityEnvironment import numpy as np +from typing import Callable, Tuple from ray.rllib.env.multi_agent_env import MultiAgentEnv from ray.rllib.utils.annotations import override +from ray.rllib.utils.types import MultiAgentDict, PolicyID, AgentID logger = logging.getLogger(__name__) @@ -26,13 +26,13 @@ class Unity3DEnv(MultiAgentEnv): """ def __init__(self, - file_name=None, - worker_id=0, - base_port=5004, - seed=0, - no_graphics=False, - timeout_wait=60, - episode_horizon=1000): + file_name: str = None, + worker_id: int = 0, + base_port: int = 5004, + seed: int = 0, + no_graphics: bool = False, + timeout_wait: int = 60, + episode_horizon: int = 1000): """Initializes a Unity3DEnv object. Args: @@ -66,6 +66,9 @@ class Unity3DEnv(MultiAgentEnv): "instead.\nMake sure you are pressing the Play (|>) button in " "your editor to start.") + import mlagents_envs + from mlagents_envs.environment import UnityEnvironment + # Try connecting to the Unity3D game instance. If a port while True: self.worker_id = worker_id @@ -92,7 +95,9 @@ class Unity3DEnv(MultiAgentEnv): self.episode_timesteps = 0 @override(MultiAgentEnv) - def step(self, action_dict): + def step( + self, action_dict: MultiAgentDict + ) -> Tuple[MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict]: """Performs one multi-agent step through the game. Args: @@ -139,7 +144,7 @@ class Unity3DEnv(MultiAgentEnv): return obs, rewards, dones, infos @override(MultiAgentEnv) - def reset(self): + def reset(self) -> MultiAgentDict: """Resets the entire Unity3D scene (a single multi-agent episode).""" self.episode_timesteps = 0 self.unity_env.reset() @@ -189,7 +194,8 @@ class Unity3DEnv(MultiAgentEnv): return obs, rewards, {"__all__": False}, infos @staticmethod - def get_policy_configs_for_game(game_name): + def get_policy_configs_for_game( + game_name: str) -> Tuple[dict, Callable[[AgentID], PolicyID]]: # The RLlib server must know about the Spaces that the Client will be # using inside Unity3D, up-front. @@ -200,7 +206,7 @@ class Unity3DEnv(MultiAgentEnv): "3DBallHard": Box(float("-inf"), float("inf"), (45, )), # SoccerStrikersVsGoalie. "Goalie": Box(float("-inf"), float("inf"), (738, )), - "Striker": Tuple([ + "Striker": TupleSpace([ Box(float("-inf"), float("inf"), (231, )), Box(float("-inf"), float("inf"), (63, )), ]), diff --git a/rllib/env/vector_env.py b/rllib/env/vector_env.py index d8da1ac6d..4267905db 100644 --- a/rllib/env/vector_env.py +++ b/rllib/env/vector_env.py @@ -1,7 +1,11 @@ import logging +import gym import numpy as np +from typing import Callable, List, Tuple from ray.rllib.utils.annotations import override, PublicAPI +from ray.rllib.utils.types import EnvType, EnvConfigDict, EnvObsType, \ + EnvInfoDict, EnvActionType logger = logging.getLogger(__name__) @@ -11,7 +15,8 @@ class VectorEnv: """An environment that supports batch evaluation using clones of sub-envs. """ - def __init__(self, observation_space, action_space, num_envs): + def __init__(self, observation_space: gym.Space, action_space: gym.Space, + num_envs: int): """Initializes a VectorEnv object. Args: @@ -25,12 +30,12 @@ class VectorEnv: self.num_envs = num_envs @staticmethod - def wrap(make_env=None, - existing_envs=None, - num_envs=1, - action_space=None, - observation_space=None, - env_config=None): + def wrap(make_env: Callable[[int], EnvType] = None, + existing_envs: List[gym.Env] = None, + num_envs: int = 1, + action_space: gym.Space = None, + observation_space: gym.Space = None, + env_config: EnvConfigDict = None): return _VectorizedGymEnv( make_env=make_env, existing_envs=existing_envs or [], @@ -40,7 +45,7 @@ class VectorEnv: env_config=env_config) @PublicAPI - def vector_reset(self): + def vector_reset(self) -> List[EnvObsType]: """Resets all sub-environments. Returns: @@ -49,7 +54,7 @@ class VectorEnv: raise NotImplementedError @PublicAPI - def reset_at(self, index): + def reset_at(self, index: int) -> EnvObsType: """Resets a single environment. Returns: @@ -58,7 +63,9 @@ class VectorEnv: raise NotImplementedError @PublicAPI - def vector_step(self, actions): + def vector_step( + self, actions: List[EnvActionType] + ) -> Tuple[List[EnvObsType], List[float], List[bool], List[EnvInfoDict]]: """Performs a vectorized step on all sub environments using `actions`. Arguments: @@ -73,7 +80,7 @@ class VectorEnv: raise NotImplementedError @PublicAPI - def get_unwrapped(self): + def get_unwrapped(self) -> List[EnvType]: """Returns the underlying sub environments. Returns: diff --git a/rllib/evaluation/episode.py b/rllib/evaluation/episode.py index 1d903dbb0..c848c7b2c 100644 --- a/rllib/evaluation/episode.py +++ b/rllib/evaluation/episode.py @@ -1,11 +1,19 @@ from collections import defaultdict import numpy as np import random +from typing import List, Dict, Callable, Any, TYPE_CHECKING from ray.rllib.env.base_env import _DUMMY_AGENT_ID +from ray.rllib.policy.policy import Policy from ray.rllib.utils import try_import_tree from ray.rllib.utils.annotations import DeveloperAPI from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray +from ray.rllib.utils.types import SampleBatchType, AgentID, PolicyID, \ + EnvObsType, EnvInfoDict, EnvActionType + +if TYPE_CHECKING: + from ray.rllib.evaluation.sample_batch_builder import \ + MultiAgentSampleBatchBuilder tree = try_import_tree() @@ -39,34 +47,42 @@ class MultiAgentEpisode: >>> episode.extra_batches.add(batch.build_and_reset()) """ - def __init__(self, policies, policy_mapping_fn, batch_builder_factory, - extra_batch_callback): - self.new_batch_builder = batch_builder_factory - self.add_extra_batch = extra_batch_callback - self.batch_builder = batch_builder_factory() - self.total_reward = 0.0 - self.length = 0 - self.episode_id = random.randrange(2e9) - self.agent_rewards = defaultdict(float) - self.custom_metrics = {} - self.user_data = {} - self.hist_data = {} - self._policies = policies - self._policy_mapping_fn = policy_mapping_fn - self._next_agent_index = 0 - self._agent_to_index = {} - self._agent_to_policy = {} - self._agent_to_rnn_state = {} - self._agent_to_last_obs = {} - self._agent_to_last_raw_obs = {} - self._agent_to_last_info = {} - self._agent_to_last_action = {} - self._agent_to_last_pi_info = {} - self._agent_to_prev_action = {} - self._agent_reward_history = defaultdict(list) + def __init__(self, policies: Dict[PolicyID, Policy], + policy_mapping_fn: Callable[[AgentID], PolicyID], + batch_builder_factory: Callable[ + [], "MultiAgentSampleBatchBuilder"], + extra_batch_callback: Callable[[SampleBatchType], None]): + self.new_batch_builder: Callable[ + [], "MultiAgentSampleBatchBuilder"] = batch_builder_factory + self.add_extra_batch: Callable[[SampleBatchType], + None] = extra_batch_callback + self.batch_builder: "MultiAgentSampleBatchBuilder" = \ + batch_builder_factory() + self.total_reward: float = 0.0 + self.length: int = 0 + self.episode_id: int = random.randrange(2e9) + self.agent_rewards: Dict[AgentID, float] = defaultdict(float) + self.custom_metrics: Dict[str, float] = {} + self.user_data: Dict[str, Any] = {} + self.hist_data: Dict[str, List[float]] = {} + self._policies: Dict[PolicyID, Policy] = policies + self._policy_mapping_fn: Callable[[AgentID], PolicyID] = \ + policy_mapping_fn + self._next_agent_index: int = 0 + self._agent_to_index: Dict[AgentID, int] = {} + self._agent_to_policy: Dict[AgentID, PolicyID] = {} + self._agent_to_rnn_state: Dict[AgentID, List[Any]] = {} + self._agent_to_last_obs: Dict[AgentID, EnvObsType] = {} + self._agent_to_last_raw_obs: Dict[AgentID, EnvObsType] = {} + self._agent_to_last_info: Dict[AgentID, EnvInfoDict] = {} + self._agent_to_last_action: Dict[AgentID, EnvActionType] = {} + self._agent_to_last_pi_info: Dict[AgentID, dict] = {} + self._agent_to_prev_action: Dict[AgentID, EnvActionType] = {} + self._agent_reward_history: Dict[AgentID, List[int]] = defaultdict( + list) @DeveloperAPI - def soft_reset(self): + def soft_reset(self) -> None: """Clears rewards and metrics, but retains RNN and other state. This is used to carry state across multiple logical episodes in the @@ -79,7 +95,7 @@ class MultiAgentEpisode: self._agent_reward_history = defaultdict(list) @DeveloperAPI - def policy_for(self, agent_id=_DUMMY_AGENT_ID): + def policy_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> Policy: """Returns the policy for the specified agent. If the agent is new, the policy mapping fn will be called to bind the @@ -91,25 +107,29 @@ class MultiAgentEpisode: return self._agent_to_policy[agent_id] @DeveloperAPI - def last_observation_for(self, agent_id=_DUMMY_AGENT_ID): + def last_observation_for( + self, agent_id: AgentID = _DUMMY_AGENT_ID) -> EnvObsType: """Returns the last observation for the specified agent.""" return self._agent_to_last_obs.get(agent_id) @DeveloperAPI - def last_raw_obs_for(self, agent_id=_DUMMY_AGENT_ID): + def last_raw_obs_for(self, + agent_id: AgentID = _DUMMY_AGENT_ID) -> EnvObsType: """Returns the last un-preprocessed obs for the specified agent.""" return self._agent_to_last_raw_obs.get(agent_id) @DeveloperAPI - def last_info_for(self, agent_id=_DUMMY_AGENT_ID): + def last_info_for(self, + agent_id: AgentID = _DUMMY_AGENT_ID) -> EnvInfoDict: """Returns the last info for the specified agent.""" return self._agent_to_last_info.get(agent_id) @DeveloperAPI - def last_action_for(self, agent_id=_DUMMY_AGENT_ID): + def last_action_for(self, + agent_id: AgentID = _DUMMY_AGENT_ID) -> EnvActionType: """Returns the last action for the specified agent, or zeros.""" if agent_id in self._agent_to_last_action: @@ -121,7 +141,8 @@ class MultiAgentEpisode: return np.zeros_like(flat) @DeveloperAPI - def prev_action_for(self, agent_id=_DUMMY_AGENT_ID): + def prev_action_for(self, + agent_id: AgentID = _DUMMY_AGENT_ID) -> EnvActionType: """Returns the previous action for the specified agent.""" if agent_id in self._agent_to_prev_action: @@ -132,7 +153,7 @@ class MultiAgentEpisode: return np.zeros_like(self.last_action_for(agent_id)) @DeveloperAPI - def prev_reward_for(self, agent_id=_DUMMY_AGENT_ID): + def prev_reward_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> float: """Returns the previous reward for the specified agent.""" history = self._agent_reward_history[agent_id] @@ -143,7 +164,7 @@ class MultiAgentEpisode: return 0.0 @DeveloperAPI - def rnn_state_for(self, agent_id=_DUMMY_AGENT_ID): + def rnn_state_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> List[Any]: """Returns the last RNN state for the specified agent.""" if agent_id not in self._agent_to_rnn_state: @@ -152,12 +173,12 @@ class MultiAgentEpisode: return self._agent_to_rnn_state[agent_id] @DeveloperAPI - def last_pi_info_for(self, agent_id=_DUMMY_AGENT_ID): + def last_pi_info_for(self, agent_id: AgentID = _DUMMY_AGENT_ID) -> dict: """Returns the last info object for the specified agent.""" return self._agent_to_last_pi_info[agent_id] - def _add_agent_rewards(self, reward_dict): + def _add_agent_rewards(self, reward_dict: Dict[AgentID, float]) -> None: for agent_id, reward in reward_dict.items(): if reward is not None: self.agent_rewards[agent_id, diff --git a/rllib/evaluation/metrics.py b/rllib/evaluation/metrics.py index 8b0442852..28a6bf460 100644 --- a/rllib/evaluation/metrics.py +++ b/rllib/evaluation/metrics.py @@ -1,6 +1,7 @@ import logging import numpy as np import collections +from typing import List, Optional, Tuple, Union import ray from ray.rllib.evaluation.rollout_metrics import RolloutMetrics @@ -8,12 +9,13 @@ from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID from ray.rllib.offline.off_policy_estimator import OffPolicyEstimate from ray.rllib.policy.policy import LEARNER_STATS_KEY from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.types import GradInfoDict, LearnerStatsDict, ResultDict logger = logging.getLogger(__name__) @DeveloperAPI -def get_learner_stats(grad_info): +def get_learner_stats(grad_info: GradInfoDict) -> LearnerStatsDict: """Return optimization stats reported from the policy. Example: @@ -36,10 +38,10 @@ def get_learner_stats(grad_info): @DeveloperAPI -def collect_metrics(local_worker=None, - remote_workers=[], - to_be_collected=[], - timeout_seconds=180): +def collect_metrics(local_worker: Optional["RolloutWorker"] = None, + remote_workers: List["ActorHandle"] = [], + to_be_collected: List["ObjectID"] = [], + timeout_seconds: int = 180) -> ResultDict: """Gathers episode metrics from RolloutWorker instances.""" episodes, to_be_collected = collect_episodes( @@ -52,10 +54,12 @@ def collect_metrics(local_worker=None, @DeveloperAPI -def collect_episodes(local_worker=None, - remote_workers=[], - to_be_collected=[], - timeout_seconds=180): +def collect_episodes( + local_worker: Optional["RolloutWorker"] = None, + remote_workers: List["ActorHandle"] = [], + to_be_collected: List["ObjectID"] = [], + timeout_seconds: int = 180 +) -> Tuple[List[Union[RolloutMetrics, OffPolicyEstimate]], List["ObjectID"]]: """Gathers new episodes metrics tuples from the given evaluators.""" if remote_workers: @@ -81,7 +85,10 @@ def collect_episodes(local_worker=None, @DeveloperAPI -def summarize_episodes(episodes, new_episodes=None): +def summarize_episodes( + episodes: List[Union[RolloutMetrics, OffPolicyEstimate]], + new_episodes: List[Union[RolloutMetrics, OffPolicyEstimate]] = None +) -> ResultDict: """Summarizes a set of episode metrics tuples. Arguments: @@ -181,7 +188,8 @@ def summarize_episodes(episodes, new_episodes=None): off_policy_estimator=dict(estimators)) -def _partition(episodes): +def _partition(episodes: List[RolloutMetrics] + ) -> Tuple[List[RolloutMetrics], List[OffPolicyEstimate]]: """Divides metrics data into true rollouts vs off-policy estimates.""" rollouts, estimates = [], [] diff --git a/rllib/evaluation/observation_function.py b/rllib/evaluation/observation_function.py index 21c29a46f..48661aa0b 100644 --- a/rllib/evaluation/observation_function.py +++ b/rllib/evaluation/observation_function.py @@ -1,9 +1,10 @@ from typing import Dict from ray.rllib.env import BaseEnv -from ray.rllib.policy import Policy, AgentID, PolicyID +from ray.rllib.policy import Policy from ray.rllib.evaluation import MultiAgentEpisode, RolloutWorker from ray.rllib.utils.framework import TensorType +from ray.rllib.utils.types import AgentID, PolicyID class ObservationFunction: diff --git a/rllib/evaluation/postprocessing.py b/rllib/evaluation/postprocessing.py index c92c27ff7..1eaa83077 100644 --- a/rllib/evaluation/postprocessing.py +++ b/rllib/evaluation/postprocessing.py @@ -4,7 +4,7 @@ from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils.annotations import DeveloperAPI -def discount(x, gamma): +def discount(x: np.ndarray, gamma: float): return scipy.signal.lfilter([1], [1, -gamma], x[::-1], axis=0)[::-1] @@ -16,12 +16,12 @@ class Postprocessing: @DeveloperAPI -def compute_advantages(rollout, - last_r, - gamma=0.9, - lambda_=1.0, - use_gae=True, - use_critic=True): +def compute_advantages(rollout: SampleBatch, + last_r: float, + gamma: float = 0.9, + lambda_: float = 1.0, + use_gae: bool = True, + use_critic: bool = True): """ Given a rollout, compute its value targets and the advantage. diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index e3f1b8c86..d61e91373 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/rllib/evaluation/rollout_worker.py @@ -5,6 +5,8 @@ import logging import pickle import platform import os +from typing import Callable, Any, List, Dict, Tuple, Union, Optional, \ + TYPE_CHECKING, TypeVar import ray from ray.util.debug import log_once, disable_log_once_globally, \ @@ -18,22 +20,35 @@ from ray.rllib.env.multi_agent_env import MultiAgentEnv from ray.rllib.env.external_multi_agent_env import ExternalMultiAgentEnv from ray.rllib.env.vector_env import VectorEnv from ray.rllib.evaluation.sampler import AsyncSampler, SyncSampler +from ray.rllib.evaluation.rollout_metrics import RolloutMetrics from ray.rllib.policy.sample_batch import MultiAgentBatch, DEFAULT_POLICY_ID from ray.rllib.policy.policy import Policy from ray.rllib.policy.tf_policy import TFPolicy from ray.rllib.policy.torch_policy import TorchPolicy from ray.rllib.offline import NoopOutput, IOContext, OutputWriter, InputReader +from ray.rllib.offline.off_policy_estimator import OffPolicyEstimator, \ + OffPolicyEstimate from ray.rllib.offline.is_estimator import ImportanceSamplingEstimator from ray.rllib.offline.wis_estimator import WeightedImportanceSamplingEstimator from ray.rllib.models import ModelCatalog -from ray.rllib.models.preprocessors import NoPreprocessor +from ray.rllib.models.preprocessors import NoPreprocessor, Preprocessor from ray.rllib.utils import merge_dicts from ray.rllib.utils.annotations import DeveloperAPI from ray.rllib.utils.debug import summarize -from ray.rllib.utils.filter import get_filter +from ray.rllib.utils.filter import get_filter, Filter from ray.rllib.utils.framework import try_import_tf, try_import_torch from ray.rllib.utils.sgd import do_minibatch_sgd from ray.rllib.utils.tf_run_builder import TFRunBuilder +from ray.rllib.utils.types import EnvType, AgentID, PolicyID, EnvConfigDict, \ + ModelConfigDict, TrainerConfigDict, SampleBatchType, ModelWeights, \ + ModelGradients, MultiAgentPolicyConfigDict + +if TYPE_CHECKING: + from ray.rllib.agents.callbacks import DefaultCallbacks + from ray.rllib.evaluation.observation_function import ObservationFunction + +# Generic type var for foreach_* methods. +T = TypeVar("T") tf = try_import_tf() torch, _ = try_import_torch() @@ -43,11 +58,11 @@ logger = logging.getLogger(__name__) # Handle to the current rollout worker, which will be set to the most recently # created RolloutWorker in this process. This can be helpful to access in # custom env or policy classes for debugging or advanced use cases. -_global_worker = None +_global_worker: "RolloutWorker" = None @DeveloperAPI -def get_global_worker(): +def get_global_worker() -> "RolloutWorker": """Returns a handle to the active rollout worker in this process.""" global _global_worker @@ -101,11 +116,11 @@ class RolloutWorker(ParallelIteratorWorker): @DeveloperAPI @classmethod def as_remote(cls, - num_cpus=None, - num_gpus=None, - memory=None, - object_store_memory=None, - resources=None): + num_cpus: int = None, + num_gpus: int = None, + memory: int = None, + object_store_memory: int = None, + resources: dict = None) -> type: return ray.remote( num_cpus=num_cpus, num_gpus=num_gpus, @@ -115,41 +130,44 @@ class RolloutWorker(ParallelIteratorWorker): @DeveloperAPI def __init__(self, - env_creator, - policy, - policy_mapping_fn=None, - policies_to_train=None, - tf_session_creator=None, - rollout_fragment_length=100, - batch_mode="truncate_episodes", - episode_horizon=None, - preprocessor_pref="deepmind", - sample_async=False, - compress_observations=False, - num_envs=1, - observation_fn=None, - observation_filter="NoFilter", - clip_rewards=None, - clip_actions=True, - env_config=None, - model_config=None, - policy_config=None, - worker_index=0, - num_workers=0, - monitor_path=None, - log_dir=None, - log_level=None, - callbacks=None, - input_creator=lambda ioctx: ioctx.default_sampler_input(), - input_evaluation=frozenset([]), - output_creator=lambda ioctx: NoopOutput(), - remote_worker_envs=False, - remote_env_batch_wait_ms=0, - soft_horizon=False, - no_done_at_end=False, - seed=None, - extra_python_environs=None, - fake_sampler=False): + env_creator: Callable[[EnvContext], EnvType], + policy: type, + policy_mapping_fn: Callable[[AgentID], PolicyID] = None, + policies_to_train: List[PolicyID] = None, + tf_session_creator: Callable[[], Any] = None, + rollout_fragment_length: int = 100, + batch_mode: str = "truncate_episodes", + episode_horizon: int = None, + preprocessor_pref: str = "deepmind", + sample_async: bool = False, + compress_observations: bool = False, + num_envs: int = 1, + observation_fn: "ObservationFunction" = None, + observation_filter: str = "NoFilter", + clip_rewards: bool = None, + clip_actions: bool = True, + env_config: EnvConfigDict = None, + model_config: ModelConfigDict = None, + policy_config: TrainerConfigDict = None, + worker_index: int = 0, + num_workers: int = 0, + monitor_path: str = None, + log_dir: str = None, + log_level: str = None, + callbacks: "DefaultCallbacks" = None, + input_creator: Callable[[ + IOContext + ], InputReader] = lambda ioctx: ioctx.default_sampler_input(), + input_evaluation: List[str] = frozenset([]), + output_creator: Callable[ + [IOContext], OutputWriter] = lambda ioctx: NoopOutput(), + remote_worker_envs: bool = False, + remote_env_batch_wait_ms: int = 0, + soft_horizon: bool = False, + no_done_at_end: bool = False, + seed: int = None, + extra_python_environs: dict = None, + fake_sampler: bool = False): """Initialize a rollout worker. Arguments: @@ -247,7 +265,7 @@ class RolloutWorker(ParallelIteratorWorker): be set. fake_sampler (bool): Use a fake (inf speed) sampler for testing. """ - self._original_kwargs = locals().copy() + self._original_kwargs: dict = locals().copy() del self._original_kwargs["self"] global _global_worker @@ -264,7 +282,7 @@ class RolloutWorker(ParallelIteratorWorker): ParallelIteratorWorker.__init__(self, gen_rollouts, False) - policy_config = policy_config or {} + policy_config: TrainerConfigDict = policy_config or {} if (tf and policy_config.get("framework") == "tfe" and not policy_config.get("no_eager_on_workers") # This eager check is necessary for certain all-framework tests @@ -281,27 +299,27 @@ class RolloutWorker(ParallelIteratorWorker): enable_periodic_logging() env_context = EnvContext(env_config or {}, worker_index) - self.policy_config = policy_config + self.policy_config: TrainerConfigDict = policy_config if callbacks: - self.callbacks = callbacks() + self.callbacks: "DefaultCallbacks" = callbacks() else: from ray.rllib.agents.callbacks import DefaultCallbacks - self.callbacks = DefaultCallbacks() - self.worker_index = worker_index - self.num_workers = num_workers - model_config = model_config or {} + self.callbacks: "DefaultCallbacks" = DefaultCallbacks() + self.worker_index: int = worker_index + self.num_workers: int = num_workers + model_config: ModelConfigDict = model_config or {} policy_mapping_fn = (policy_mapping_fn or (lambda agent_id: DEFAULT_POLICY_ID)) if not callable(policy_mapping_fn): raise ValueError("Policy mapping function not callable?") - self.env_creator = env_creator - self.rollout_fragment_length = rollout_fragment_length * num_envs - self.batch_mode = batch_mode - self.compress_observations = compress_observations - self.preprocessing_enabled = True - self.last_batch = None - self.global_vars = None - self.fake_sampler = fake_sampler + self.env_creator: Callable[[EnvContext], EnvType] = env_creator + self.rollout_fragment_length: int = rollout_fragment_length * num_envs + self.batch_mode: str = batch_mode + self.compress_observations: bool = compress_observations + self.preprocessing_enabled: bool = True + self.last_batch: SampleBatchType = None + self.global_vars: dict = None + self.fake_sampler: bool = fake_sampler self.env = _validate_env(env_creator(env_context)) if isinstance(self.env, (BaseEnv, MultiAgentEnv)): @@ -336,7 +354,7 @@ class RolloutWorker(ParallelIteratorWorker): env = wrappers.Monitor(env, monitor_path, resume=True) return env - self.env = wrap(self.env) + self.env: EnvType = wrap(self.env) def make_env(vector_index): return wrap( @@ -346,7 +364,11 @@ class RolloutWorker(ParallelIteratorWorker): self.tf_sess = None policy_dict = _validate_and_canonicalize(policy, self.env) - self.policies_to_train = policies_to_train or list(policy_dict.keys()) + self.policies_to_train: List[PolicyID] = policies_to_train or list( + policy_dict.keys()) + self.policy_map: Dict[PolicyID, Policy] = None + self.preprocessors: Dict[PolicyID, Preprocessor] = None + # set numpy and python seed if seed is not None: np.random.seed(seed) @@ -393,7 +415,8 @@ class RolloutWorker(ParallelIteratorWorker): self.policy_map, self.preprocessors = self._build_policy_map( policy_dict, policy_config) - self.multiagent = set(self.policy_map.keys()) != {DEFAULT_POLICY_ID} + self.multiagent: bool = set( + self.policy_map.keys()) != {DEFAULT_POLICY_ID} if self.multiagent: if not ((isinstance(self.env, MultiAgentEnv) or isinstance(self.env, ExternalMultiAgentEnv)) @@ -404,7 +427,7 @@ class RolloutWorker(ParallelIteratorWorker): "{} is not a subclass of BaseEnv, MultiAgentEnv or " "ExternalMultiAgentEnv?".format(self.env)) - self.filters = { + self.filters: Dict[PolicyID, Filter] = { policy_id: get_filter(observation_filter, policy.observation_space.shape) for (policy_id, policy) in self.policy_map.items() @@ -413,13 +436,13 @@ class RolloutWorker(ParallelIteratorWorker): logger.info("Built filter map: {}".format(self.filters)) # Always use vector env for consistency even if num_envs = 1. - self.async_env = BaseEnv.to_base_env( + self.async_env: BaseEnv = BaseEnv.to_base_env( self.env, make_env=make_env, num_envs=num_envs, remote_envs=remote_worker_envs, remote_env_batch_wait_ms=remote_env_batch_wait_ms) - self.num_envs = num_envs + self.num_envs: int = num_envs # `truncate_episodes`: Allow a batch to contain more than one episode # (fragments) and always make the batch `rollout_fragment_length` @@ -435,8 +458,9 @@ class RolloutWorker(ParallelIteratorWorker): raise ValueError("Unsupported batch mode: {}".format( self.batch_mode)) - self.io_context = IOContext(log_dir, policy_config, worker_index, self) - self.reward_estimators = [] + self.io_context: IOContext = IOContext(log_dir, policy_config, + worker_index, self) + self.reward_estimators: OffPolicyEstimator = [] for method in input_evaluation: if method == "simulation": logger.warning( @@ -494,24 +518,21 @@ class RolloutWorker(ParallelIteratorWorker): no_done_at_end=no_done_at_end, observation_fn=observation_fn) - self.input_reader = input_creator(self.io_context) - assert isinstance(self.input_reader, InputReader), self.input_reader - self.output_writer = output_creator(self.io_context) - assert isinstance(self.output_writer, OutputWriter), self.output_writer + self.input_reader: InputReader = input_creator(self.io_context) + self.output_writer: OutputWriter = output_creator(self.io_context) logger.debug( "Created rollout worker with env {} ({}), policies {}".format( self.async_env, self.env, self.policy_map)) @DeveloperAPI - def sample(self): + def sample(self) -> SampleBatchType: """Returns a batch of experience sampled from this worker. This method must be implemented by subclasses. Returns: - Union[SampleBatch,MultiAgentBatch]: A columnar batch of experiences - (e.g., tensors), or a multi-agent batch. + SampleBatchType: A columnar batch of experiences (e.g., tensors). Examples: >>> print(worker.sample()) @@ -569,13 +590,14 @@ class RolloutWorker(ParallelIteratorWorker): @DeveloperAPI @ray.method(num_return_vals=2) - def sample_with_count(self): + def sample_with_count(self) -> Tuple[SampleBatchType, int]: """Same as sample() but returns the count as a separate future.""" batch = self.sample() return batch, batch.count @DeveloperAPI - def get_weights(self, policies=None): + def get_weights(self, + policies: List[PolicyID] = None) -> (ModelWeights, dict): """Returns the model weights of this worker. Returns: @@ -593,7 +615,8 @@ class RolloutWorker(ParallelIteratorWorker): } @DeveloperAPI - def set_weights(self, weights, global_vars=None): + def set_weights(self, weights: ModelWeights, + global_vars: dict = None) -> None: """Sets the model weights of this worker. Examples: @@ -606,7 +629,8 @@ class RolloutWorker(ParallelIteratorWorker): self.set_global_vars(global_vars) @DeveloperAPI - def compute_gradients(self, samples): + def compute_gradients( + self, samples: SampleBatchType) -> Tuple[ModelGradients, dict]: """Returns a gradient computed w.r.t the specified samples. Returns: @@ -650,7 +674,7 @@ class RolloutWorker(ParallelIteratorWorker): return grad_out, info_out @DeveloperAPI - def apply_gradients(self, grads): + def apply_gradients(self, grads: ModelGradients) -> Dict[PolicyID, Any]: """Applies the given gradients to this worker's weights. Examples: @@ -678,7 +702,7 @@ class RolloutWorker(ParallelIteratorWorker): return self.policy_map[DEFAULT_POLICY_ID].apply_gradients(grads) @DeveloperAPI - def learn_on_batch(self, samples): + def learn_on_batch(self, samples: SampleBatchType) -> dict: """Update policies based on the given batch. This is the equivalent to apply_gradients(compute_gradients(samples)), @@ -721,8 +745,9 @@ class RolloutWorker(ParallelIteratorWorker): logger.debug("Training out:\n\n{}\n".format(summarize(info_out))) return info_out - def sample_and_learn(self, expected_batch_size, num_sgd_iter, - sgd_minibatch_size, standardize_fields): + def sample_and_learn(self, expected_batch_size: int, num_sgd_iter: int, + sgd_minibatch_size: str, + standardize_fields: List[str]) -> Tuple[dict, int]: """Sample and batch and learn on it. This is typically used in combination with distributed allreduce. @@ -749,7 +774,7 @@ class RolloutWorker(ParallelIteratorWorker): return info, batch.count @DeveloperAPI - def get_metrics(self): + def get_metrics(self) -> List[Union[RolloutMetrics, OffPolicyEstimate]]: """Returns a list of new RolloutMetric objects from evaluation.""" out = self.sampler.get_metrics() @@ -758,7 +783,7 @@ class RolloutWorker(ParallelIteratorWorker): return out @DeveloperAPI - def foreach_env(self, func): + def foreach_env(self, func: Callable[[BaseEnv], T]) -> List[T]: """Apply the given function to each underlying env instance.""" envs = self.async_env.get_unwrapped() @@ -768,7 +793,8 @@ class RolloutWorker(ParallelIteratorWorker): return [func(e) for e in envs] @DeveloperAPI - def get_policy(self, policy_id=DEFAULT_POLICY_ID): + def get_policy( + self, policy_id: Optional[PolicyID] = DEFAULT_POLICY_ID) -> Policy: """Return policy for the specified id, or None. Arguments: @@ -778,19 +804,22 @@ class RolloutWorker(ParallelIteratorWorker): return self.policy_map.get(policy_id) @DeveloperAPI - def for_policy(self, func, policy_id=DEFAULT_POLICY_ID): + def for_policy(self, + func: Callable[[Policy], T], + policy_id: Optional[PolicyID] = DEFAULT_POLICY_ID) -> T: """Apply the given function to the specified policy.""" return func(self.policy_map[policy_id]) @DeveloperAPI - def foreach_policy(self, func): + def foreach_policy(self, func: Callable[[Policy, PolicyID], T]) -> List[T]: """Apply the given function to each (policy, policy_id) tuple.""" return [func(policy, pid) for pid, policy in self.policy_map.items()] @DeveloperAPI - def foreach_trainable_policy(self, func): + def foreach_trainable_policy( + self, func: Callable[[Policy, PolicyID], T]) -> List[T]: """ Applies the given function to each (policy, policy_id) tuple, which can be found in `self.policies_to_train`. @@ -809,7 +838,7 @@ class RolloutWorker(ParallelIteratorWorker): ] @DeveloperAPI - def sync_filters(self, new_filters): + def sync_filters(self, new_filters: dict) -> None: """Changes self's filter to given and rebases any accumulated delta. Args: @@ -820,7 +849,7 @@ class RolloutWorker(ParallelIteratorWorker): self.filters[k].sync(new_filters[k]) @DeveloperAPI - def get_filters(self, flush_after=False): + def get_filters(self, flush_after: bool = False) -> dict: """Returns a snapshot of filters. Args: @@ -837,7 +866,7 @@ class RolloutWorker(ParallelIteratorWorker): return return_filters @DeveloperAPI - def save(self): + def save(self) -> str: filters = self.get_filters(flush_after=True) state = { pid: self.policy_map[pid].get_state() @@ -846,61 +875,66 @@ class RolloutWorker(ParallelIteratorWorker): return pickle.dumps({"filters": filters, "state": state}) @DeveloperAPI - def restore(self, objs): + def restore(self, objs: str) -> None: objs = pickle.loads(objs) self.sync_filters(objs["filters"]) for pid, state in objs["state"].items(): self.policy_map[pid].set_state(state) @DeveloperAPI - def set_global_vars(self, global_vars): + def set_global_vars(self, global_vars: dict) -> None: self.foreach_policy(lambda p, _: p.on_global_var_update(global_vars)) self.global_vars = global_vars @DeveloperAPI - def get_global_vars(self): + def get_global_vars(self) -> dict: return self.global_vars @DeveloperAPI - def export_policy_model(self, export_dir, policy_id=DEFAULT_POLICY_ID): + def export_policy_model(self, + export_dir: str, + policy_id: PolicyID = DEFAULT_POLICY_ID): self.policy_map[policy_id].export_model(export_dir) @DeveloperAPI def import_policy_model_from_h5(self, - import_file, - policy_id=DEFAULT_POLICY_ID): + import_file: str, + policy_id: PolicyID = DEFAULT_POLICY_ID): self.policy_map[policy_id].import_model_from_h5(import_file) @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): self.policy_map[policy_id].export_checkpoint(export_dir, filename_prefix) @DeveloperAPI - def stop(self): + def stop(self) -> None: self.async_env.stop() @DeveloperAPI - def creation_args(self): + def creation_args(self) -> dict: """Returns the args used to create this worker.""" return self._original_kwargs @DeveloperAPI - def get_host(self): + def get_host(self) -> str: """Returns the hostname of the process running this evaluator.""" return platform.node() @DeveloperAPI - def apply(self, func, *args): - """Apply the given function to this evaluator instance.""" + def apply(self, func: Callable[["RolloutWorker"], T], *args) -> T: + """Apply the given function to this rollout worker instance.""" return func(self, *args) - def _build_policy_map(self, policy_dict, policy_config): + def _build_policy_map( + self, policy_dict: MultiAgentPolicyConfigDict, + policy_config: TrainerConfigDict + ) -> Tuple[Dict[PolicyID, Policy], Dict[PolicyID, Preprocessor]]: policy_map = {} preprocessors = {} for name, (cls, obs_space, act_space, @@ -942,7 +976,8 @@ class RolloutWorker(ParallelIteratorWorker): logger.info("Built preprocessor map: {}".format(preprocessors)) return policy_map, preprocessors - def setup_torch_data_parallel(self, url, world_rank, world_size, backend): + def setup_torch_data_parallel(self, url: str, world_rank: int, + world_size: int, backend: str) -> None: """Join a torch process group for distributed SGD.""" logger.info("Joining process group, url={}, world_rank={}, " @@ -960,11 +995,11 @@ class RolloutWorker(ParallelIteratorWorker): "This policy does not support torch distributed", policy) policy.distributed_world_size = world_size - def get_node_ip(self): + def get_node_ip(self) -> str: """Returns the IP address of the current node.""" return ray.services.get_node_ip_address() - def find_free_port(self): + def find_free_port(self) -> int: """Finds a free port on the current node.""" from ray.util.sgd import utils return utils.find_free_port() @@ -974,7 +1009,8 @@ class RolloutWorker(ParallelIteratorWorker): self.sampler.shutdown = True -def _validate_and_canonicalize(policy, env): +def _validate_and_canonicalize(policy: Policy, + env: EnvType) -> MultiAgentPolicyConfigDict: if isinstance(policy, dict): _validate_multiagent_config(policy) return policy @@ -992,7 +1028,8 @@ def _validate_and_canonicalize(policy, env): } -def _validate_multiagent_config(policy, allow_none_graph=False): +def _validate_multiagent_config(policy: MultiAgentPolicyConfigDict, + allow_none_graph: bool = False): for k, v in policy.items(): if not isinstance(k, str): raise ValueError("policy keys must be strs, got {}".format( @@ -1019,7 +1056,7 @@ def _validate_multiagent_config(policy, allow_none_graph=False): "got {}".format(type(v[3]))) -def _validate_env(env): +def _validate_env(env: Any) -> EnvType: # allow this as a special case (assumed gym.Env) if hasattr(env, "observation_space") and hasattr(env, "action_space"): return env @@ -1033,7 +1070,7 @@ def _validate_env(env): return env -def _has_tensorflow_graph(policy_dict): +def _has_tensorflow_graph(policy_dict: MultiAgentPolicyConfigDict) -> bool: for policy, _, _, _ in policy_dict.values(): if issubclass(policy, TFPolicy): return True diff --git a/rllib/evaluation/sample_batch_builder.py b/rllib/evaluation/sample_batch_builder.py index e99844883..f5e3b002f 100644 --- a/rllib/evaluation/sample_batch_builder.py +++ b/rllib/evaluation/sample_batch_builder.py @@ -1,17 +1,24 @@ import collections import logging import numpy as np +from typing import List, Any, Dict, Optional, TYPE_CHECKING +from ray.rllib.evaluation.episode import MultiAgentEpisode +from ray.rllib.policy.policy import Policy from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch from ray.rllib.utils.annotations import PublicAPI, DeveloperAPI from ray.rllib.utils.debug import summarize +from ray.rllib.utils.types import PolicyID, AgentID from ray.rllib.env.base_env import _DUMMY_AGENT_ID from ray.util.debug import log_once +if TYPE_CHECKING: + from ray.rllib.agents.callbacks import DefaultCallbacks + logger = logging.getLogger(__name__) -def to_float_array(v): +def to_float_array(v: List[Any]) -> np.ndarray: arr = np.array(v) if arr.dtype == np.float64: return arr.astype(np.float32) # save some memory @@ -30,11 +37,11 @@ class SampleBatchBuilder: @PublicAPI def __init__(self): - self.buffers = collections.defaultdict(list) + self.buffers: Dict[str, List] = collections.defaultdict(list) self.count = 0 @PublicAPI - def add_values(self, **values): + def add_values(self, **values: Dict[str, Any]) -> None: """Add the given dictionary (row) of values to this batch.""" for k, v in values.items(): @@ -42,7 +49,7 @@ class SampleBatchBuilder: self.count += 1 @PublicAPI - def add_batch(self, batch): + def add_batch(self, batch: SampleBatch) -> None: """Add the given batch of values to this batch.""" for k, column in batch.items(): @@ -50,7 +57,7 @@ class SampleBatchBuilder: self.count += batch.count @PublicAPI - def build_and_reset(self): + def build_and_reset(self) -> SampleBatch: """Returns a sample batch including all previously added values.""" batch = SampleBatch( @@ -75,7 +82,8 @@ class MultiAgentSampleBatchBuilder: corresponding policy batch for the agent's policy. """ - def __init__(self, policy_map, clip_rewards, callbacks): + def __init__(self, policy_map: Dict[PolicyID, Policy], clip_rewards: bool, + callbacks: "DefaultCallbacks"): """Initialize a MultiAgentSampleBatchBuilder. Args: @@ -102,7 +110,7 @@ class MultiAgentSampleBatchBuilder: # Regardless of the number of agents involved in each of these steps. self.count = 0 - def total(self): + def total(self) -> int: """Returns the total number of steps taken in the env (all agents). Returns: @@ -112,7 +120,7 @@ class MultiAgentSampleBatchBuilder: return sum(a.count for a in self.agent_builders.values()) - def has_pending_agent_data(self): + def has_pending_agent_data(self) -> bool: """Returns whether there is pending unprocessed data. Returns: @@ -123,7 +131,8 @@ class MultiAgentSampleBatchBuilder: return len(self.agent_builders) > 0 @DeveloperAPI - def add_values(self, agent_id, policy_id, **values): + def add_values(self, agent_id: AgentID, policy_id: AgentID, + **values: Dict[str, Any]) -> None: """Add the given dictionary (row) of values to this batch. Arguments: @@ -142,7 +151,8 @@ class MultiAgentSampleBatchBuilder: self.agent_builders[agent_id].add_values(**values) - def postprocess_batch_so_far(self, episode=None): + def postprocess_batch_so_far( + self, episode: Optional[MultiAgentEpisode] = None) -> None: """Apply policy postprocessors to any unprocessed rows. This pushes the postprocessed per-agent batches onto the per-policy @@ -210,7 +220,7 @@ class MultiAgentSampleBatchBuilder: self.agent_builders.clear() self.agent_to_policy.clear() - def check_missing_dones(self): + def check_missing_dones(self) -> None: for agent_id, builder in self.agent_builders.items(): if builder.buffers["dones"][-1] is not True: raise ValueError( @@ -223,7 +233,8 @@ class MultiAgentSampleBatchBuilder: "Alternatively, set no_done_at_end=True to allow this.") @DeveloperAPI - def build_and_reset(self, episode=None): + def build_and_reset(self, episode: Optional[MultiAgentEpisode] = None + ) -> MultiAgentBatch: """Returns the accumulated sample batches for each policy. Any unprocessed rows will be first postprocessed with a policy diff --git a/rllib/evaluation/sampler.py b/rllib/evaluation/sampler.py index 42e6d9f19..614f5beb5 100644 --- a/rllib/evaluation/sampler.py +++ b/rllib/evaluation/sampler.py @@ -5,14 +5,18 @@ import numpy as np import queue import threading import time +from typing import List, Dict, Callable, Set, Tuple, Any, Iterable, Union, \ + TYPE_CHECKING from ray.util.debug import log_once from ray.rllib.evaluation.episode import MultiAgentEpisode from ray.rllib.evaluation.rollout_metrics import RolloutMetrics from ray.rllib.evaluation.sample_batch_builder import \ MultiAgentSampleBatchBuilder -from ray.rllib.policy.policy import clip_action +from ray.rllib.policy.policy import clip_action, Policy from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.models.preprocessors import Preprocessor +from ray.rllib.utils.filter import Filter from ray.rllib.env.base_env import BaseEnv, ASYNC_RESET_RETURN from ray.rllib.env.atari_wrappers import get_wrapper_by_cls, MonitorEnv from ray.rllib.offline import InputReader @@ -22,6 +26,14 @@ from ray.rllib.utils.debug import summarize from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray, \ unbatch from ray.rllib.utils.tf_run_builder import TFRunBuilder +from ray.rllib.utils.types import SampleBatchType, AgentID, PolicyID, \ + EnvObsType, EnvInfoDict, EnvID, MultiEnvDict, EnvActionType, \ + TensorStructType + +if TYPE_CHECKING: + from ray.rllib.agents.callbacks import DefaultCallbacks + from ray.rllib.evaluation.observation_function import ObservationFunction + from ray.rllib.evaluation.rollout_worker import RolloutWorker tree = try_import_tree() @@ -32,8 +44,11 @@ PolicyEvalData = namedtuple("PolicyEvalData", [ "prev_reward" ]) +# A batch of RNN states with dimensions [state_index, batch, state_object]. +StateBatch = List[List[Any]] -class PerfStats: + +class _PerfStats: """Sampler perf stats that will be included in rollout metrics.""" def __init__(self): @@ -55,7 +70,7 @@ class SamplerInput(InputReader, metaclass=ABCMeta): """Reads input experiences from an existing sampler.""" @override(InputReader) - def next(self): + def next(self) -> SampleBatchType: batches = [self.get_data()] batches.extend(self.get_extra_batches()) if len(batches) > 1: @@ -65,17 +80,17 @@ class SamplerInput(InputReader, metaclass=ABCMeta): @abstractmethod @DeveloperAPI - def get_data(self): + def get_data(self) -> SampleBatchType: raise NotImplementedError @abstractmethod @DeveloperAPI - def get_metrics(self): + def get_metrics(self) -> List[RolloutMetrics]: raise NotImplementedError @abstractmethod @DeveloperAPI - def get_extra_batches(self): + def get_extra_batches(self) -> List[SampleBatchType]: raise NotImplementedError @@ -86,22 +101,22 @@ class SyncSampler(SamplerInput): def __init__(self, *, - worker, - env, - policies, - policy_mapping_fn, - preprocessors, - obs_filters, - clip_rewards, - rollout_fragment_length, - callbacks, - horizon=None, - pack_multiple_episodes_in_batch=False, + worker: "RolloutWorker", + env: BaseEnv, + policies: Dict[PolicyID, Policy], + policy_mapping_fn: Callable[[AgentID], PolicyID], + preprocessors: Dict[PolicyID, Preprocessor], + obs_filters: Dict[PolicyID, Filter], + clip_rewards: bool, + rollout_fragment_length: int, + callbacks: "DefaultCallbacks", + horizon: int = None, + pack_multiple_episodes_in_batch: bool = False, tf_sess=None, - clip_actions=True, - soft_horizon=False, - no_done_at_end=False, - observation_fn=None): + clip_actions: bool = True, + soft_horizon: bool = False, + no_done_at_end: bool = False, + observation_fn: "ObservationFunction" = None): """Initializes a SyncSampler object. Args: @@ -148,7 +163,7 @@ class SyncSampler(SamplerInput): self.preprocessors = preprocessors self.obs_filters = obs_filters self.extra_batches = queue.Queue() - self.perf_stats = PerfStats() + self.perf_stats = _PerfStats() # Create the rollout generator to use for calls to `get_data()`. self.rollout_provider = _env_runner( worker, self.base_env, self.extra_batches.put, self.policies, @@ -159,7 +174,7 @@ class SyncSampler(SamplerInput): self.metrics_queue = queue.Queue() @override(SamplerInput) - def get_data(self): + def get_data(self) -> SampleBatchType: while True: item = next(self.rollout_provider) if isinstance(item, RolloutMetrics): @@ -168,7 +183,7 @@ class SyncSampler(SamplerInput): return item @override(SamplerInput) - def get_metrics(self): + def get_metrics(self) -> List[RolloutMetrics]: completed = [] while True: try: @@ -179,7 +194,7 @@ class SyncSampler(SamplerInput): return completed @override(SamplerInput) - def get_extra_batches(self): + def get_extra_batches(self) -> List[SampleBatchType]: extra = [] while True: try: @@ -199,37 +214,37 @@ class AsyncSampler(threading.Thread, SamplerInput): def __init__(self, *, - worker, - env, - policies, - policy_mapping_fn, - preprocessors, - obs_filters, - clip_rewards, - rollout_fragment_length, - callbacks, - horizon=None, - pack_multiple_episodes_in_batch=False, + worker: "RolloutWorker", + env: BaseEnv, + policies: Dict[PolicyID, Policy], + policy_mapping_fn: Callable[[AgentID], PolicyID], + preprocessors: Dict[PolicyID, Preprocessor], + obs_filters: Dict[PolicyID, Filter], + clip_rewards: bool, + rollout_fragment_length: int, + callbacks: "DefaultCallbacks", + horizon: int = None, + pack_multiple_episodes_in_batch: bool = False, tf_sess=None, - clip_actions=True, - blackhole_outputs=False, - soft_horizon=False, - no_done_at_end=False, - observation_fn=None): + clip_actions: bool = True, + blackhole_outputs: bool = False, + soft_horizon: bool = False, + no_done_at_end: bool = False, + observation_fn: "ObservationFunction" = None): """Initializes a AsyncSampler object. Args: worker (RolloutWorker): The RolloutWorker that will use this Sampler for sampling. env (Env): Any Env object. Will be converted into an RLlib BaseEnv. - policies (Dict[str,Policy]): Mapping from policy ID to Policy obj. + policies (Dict[str, Policy]): Mapping from policy ID to Policy obj. policy_mapping_fn (callable): Callable that takes an agent ID and returns a Policy object. - preprocessors (Dict[str,Preprocessor]): Mapping from policy ID to + preprocessors (Dict[str, Preprocessor]): Mapping from policy ID to Preprocessor object for the observations prior to filtering. - obs_filters (Dict[str,Filter]): Mapping from policy ID to + obs_filters (Dict[str, Filter]): Mapping from policy ID to env Filter object. - clip_rewards (Union[bool,float]): True for +/-1.0 clipping, actual + clip_rewards (Union[bool, float]): True for +/-1.0 clipping, actual float value for +/- value clipping. False for no clipping. rollout_fragment_length (int): The length of a fragment to collect before building a SampleBatch from the data and resetting @@ -279,7 +294,7 @@ class AsyncSampler(threading.Thread, SamplerInput): self.blackhole_outputs = blackhole_outputs self.soft_horizon = soft_horizon self.no_done_at_end = no_done_at_end - self.perf_stats = PerfStats() + self.perf_stats = _PerfStats() self.shutdown = False self.observation_fn = observation_fn @@ -317,7 +332,7 @@ class AsyncSampler(threading.Thread, SamplerInput): queue_putter(item) @override(SamplerInput) - def get_data(self): + def get_data(self) -> SampleBatchType: if not self.is_alive(): raise RuntimeError("Sampling thread has died") rollout = self.queue.get(timeout=600.0) @@ -329,7 +344,7 @@ class AsyncSampler(threading.Thread, SamplerInput): return rollout @override(SamplerInput) - def get_metrics(self): + def get_metrics(self) -> List[RolloutMetrics]: completed = [] while True: try: @@ -340,7 +355,7 @@ class AsyncSampler(threading.Thread, SamplerInput): return completed @override(SamplerInput) - def get_extra_batches(self): + def get_extra_batches(self) -> List[SampleBatchType]: extra = [] while True: try: @@ -350,11 +365,17 @@ class AsyncSampler(threading.Thread, SamplerInput): return extra -def _env_runner(worker, base_env, extra_batch_callback, policies, - policy_mapping_fn, rollout_fragment_length, horizon, - preprocessors, obs_filters, clip_rewards, clip_actions, - pack_multiple_episodes_in_batch, callbacks, tf_sess, - perf_stats, soft_horizon, no_done_at_end, observation_fn): +def _env_runner( + worker: "RolloutWorker", base_env: BaseEnv, + extra_batch_callback: Callable[[SampleBatchType], None], policies, + policy_mapping_fn: Callable[[AgentID], PolicyID], + rollout_fragment_length: int, horizon: int, + preprocessors: Dict[PolicyID, Preprocessor], + obs_filters: Dict[PolicyID, Filter], clip_rewards: bool, + clip_actions: bool, pack_multiple_episodes_in_batch: bool, + callbacks: "DefaultCallbacks", tf_sess, perf_stats: _PerfStats, + soft_horizon: bool, no_done_at_end: bool, + observation_fn: "ObservationFunction") -> Iterable[SampleBatchType]: """This implements the common experience collection logic. Args: @@ -381,7 +402,7 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, callbacks (DefaultCallbacks): User callbacks to run on episode events. tf_sess (Session|None): Optional tensorflow session to use for batching TF policy evaluations. - perf_stats (PerfStats): Record perf stats into this object. + perf_stats (_PerfStats): Record perf stats into this object. soft_horizon (bool): Calculate rewards but don't reset the environment when the horizon is hit. no_done_at_end (bool): Ignore the done=True at the end of the episode @@ -424,7 +445,7 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, # Pool of batch builders, which can be shared across episodes to pack # trajectory data. - batch_builder_pool = [] + batch_builder_pool: List[MultiAgentSampleBatchBuilder] = [] def get_batch_builder(): if batch_builder_pool: @@ -437,6 +458,7 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, episode = MultiAgentEpisode(policies, policy_mapping_fn, get_batch_builder, extra_batch_callback) # Call each policy's Exploration.on_episode_start method. + # type: Policy for p in policies.values(): if getattr(p, "exploration", None) is not None: p.exploration.on_episode_start( @@ -451,12 +473,13 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, episode=episode) return episode - active_episodes = defaultdict(new_episode) + active_episodes: Dict[str, MultiAgentEpisode] = defaultdict(new_episode) while True: perf_stats.iters += 1 t0 = time.time() # Get observations from all ready agents. + # type: MultiEnvDict, MultiEnvDict, MultiEnvDict, MultiEnvDict, ... unfiltered_obs, rewards, dones, infos, off_policy_actions = \ base_env.poll() perf_stats.env_wait_time += time.time() - t0 @@ -468,6 +491,8 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, # Process observations and prepare for policy evaluation. t1 = time.time() + # type: Set[EnvID], Dict[PolicyID, List[PolicyEvalData]], + # List[Union[RolloutMetrics, SampleBatchType]] active_envs, to_eval, outputs = _process_observations( worker=worker, base_env=base_env, @@ -493,6 +518,7 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, # Do batched policy eval (accross vectorized envs). t2 = time.time() + # type: Dict[PolicyID, Tuple[TensorStructType, StateBatch, dict]] eval_results = _do_policy_eval( to_eval=to_eval, policies=policies, @@ -502,14 +528,15 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, # Process results and update episode state. t3 = time.time() - actions_to_send = _process_policy_eval_results( - to_eval=to_eval, - eval_results=eval_results, - active_episodes=active_episodes, - active_envs=active_envs, - off_policy_actions=off_policy_actions, - policies=policies, - clip_actions=clip_actions) + actions_to_send: Dict[EnvID, Dict[AgentID, EnvActionType]] = \ + _process_policy_eval_results( + to_eval=to_eval, + eval_results=eval_results, + active_episodes=active_episodes, + active_envs=active_envs, + off_policy_actions=off_policy_actions, + policies=policies, + clip_actions=clip_actions) perf_stats.processing_time += time.time() - t3 # Return computed actions to ready envs. We also send to envs that have @@ -520,10 +547,21 @@ def _env_runner(worker, base_env, extra_batch_callback, policies, def _process_observations( - worker, base_env, policies, batch_builder_pool, active_episodes, - unfiltered_obs, rewards, dones, infos, horizon, preprocessors, - obs_filters, rollout_fragment_length, pack_multiple_episodes_in_batch, - callbacks, soft_horizon, no_done_at_end, observation_fn): + worker: "RolloutWorker", base_env: BaseEnv, + policies: Dict[PolicyID, Policy], + batch_builder_pool: List[MultiAgentSampleBatchBuilder], + active_episodes: Dict[str, MultiAgentEpisode], + unfiltered_obs: Dict[EnvID, Dict[AgentID, EnvObsType]], + rewards: Dict[EnvID, Dict[AgentID, float]], + dones: Dict[EnvID, Dict[AgentID, bool]], + infos: Dict[EnvID, Dict[AgentID, EnvInfoDict]], horizon: int, + preprocessors: Dict[PolicyID, Preprocessor], + obs_filters: Dict[PolicyID, Filter], rollout_fragment_length: int, + pack_multiple_episodes_in_batch: bool, callbacks: "DefaultCallbacks", + soft_horizon: bool, no_done_at_end: bool, + observation_fn: "ObservationFunction" +) -> Tuple[Set[EnvID], Dict[PolicyID, List[PolicyEvalData]], List[Union[ + RolloutMetrics, SampleBatchType]]]: """Record new data from the environment and prepare for policy evaluation. Args: @@ -532,7 +570,7 @@ def _process_observations( policies (dict): Map of policy ids to Policy instances. batch_builder_pool (List[SampleBatchBuilder]): List of pooled SampleBatchBuilder object for recycling. - active_episodes (defaultdict[str,MultiAgentEpisode]): Mapping from + active_episodes (Dict[str, MultiAgentEpisode]): Mapping from episode ID to currently ongoing MultiAgentEpisode object. unfiltered_obs (dict): Doubly keyed dict of env-ids -> agent ids -> unfiltered observation tensor, returned by a `BaseEnv.poll()` call. @@ -568,16 +606,18 @@ def _process_observations( - outputs: List of metrics and samples to return from the sampler. """ - active_envs = set() - to_eval = defaultdict(list) - outputs = [] - large_batch_threshold = max(1000, rollout_fragment_length * 10) if \ + # Output objects. + active_envs: Set[EnvID] = set() + to_eval: Dict[PolicyID, List[PolicyEvalData]] = defaultdict(list) + outputs: List[Union[RolloutMetrics, SampleBatchType]] = [] + + large_batch_threshold: int = max(1000, rollout_fragment_length * 10) if \ rollout_fragment_length != float("inf") else 5000 - # For each environment. + # type: EnvID, Dict[AgentID, EnvObsType] for env_id, agent_obs in unfiltered_obs.items(): - is_new_episode = env_id not in active_episodes - episode = active_episodes[env_id] + is_new_episode: bool = env_id not in active_episodes + episode: MultiAgentEpisode = active_episodes[env_id] if not is_new_episode: episode.length += 1 episode.batch_builder.count += 1 @@ -605,7 +645,8 @@ def _process_observations( hit_horizon = (episode.length >= horizon and not dones[env_id]["__all__"]) all_agents_done = True - atari_metrics = _fetch_atari_metrics(base_env) + atari_metrics: List[RolloutMetrics] = _fetch_atari_metrics( + base_env) if atari_metrics is not None: for m in atari_metrics: outputs.append( @@ -623,7 +664,7 @@ def _process_observations( # Custom observation function is applied before preprocessing. if observation_fn: - agent_obs = observation_fn( + agent_obs: Dict[AgentID, EnvObsType] = observation_fn( agent_obs=agent_obs, worker=worker, base_env=base_env, @@ -634,15 +675,17 @@ def _process_observations( "observe() must return a dict of agent observations") # For each agent in the environment. + # type: AgentID, EnvObsType for agent_id, raw_obs in agent_obs.items(): assert agent_id != "__all__" - policy_id = episode.policy_for(agent_id) - prep_obs = _get_or_raise(preprocessors, - policy_id).transform(raw_obs) + policy_id: PolicyID = episode.policy_for(agent_id) + prep_obs: EnvObsType = _get_or_raise(preprocessors, + policy_id).transform(raw_obs) if log_once("prep_obs"): logger.info("Preprocessed obs: {}".format(summarize(prep_obs))) - filtered_obs = _get_or_raise(obs_filters, policy_id)(prep_obs) + filtered_obs: EnvObsType = _get_or_raise(obs_filters, + policy_id)(prep_obs) if log_once("filtered_obs"): logger.info("Filtered obs: {}".format(summarize(filtered_obs))) @@ -655,7 +698,8 @@ def _process_observations( episode.last_action_for(agent_id), rewards[env_id][agent_id] or 0.0)) - last_observation = episode.last_observation_for(agent_id) + last_observation: EnvObsType = episode.last_observation_for( + agent_id) episode._set_last_observation(agent_id, filtered_obs) episode._set_last_raw_obs(agent_id, raw_obs) episode._set_last_info(agent_id, infos[env_id].get(agent_id, {})) @@ -722,10 +766,11 @@ def _process_observations( episode=episode) if hit_horizon and soft_horizon: episode.soft_reset() - resetted_obs = agent_obs + resetted_obs: Dict[AgentID, EnvObsType] = agent_obs else: del active_episodes[env_id] - resetted_obs = base_env.try_reset(env_id) + resetted_obs: Dict[AgentID, EnvObsType] = base_env.try_reset( + env_id) if resetted_obs is None: # Reset not supported, drop this env from the ready list. if horizon != float("inf"): @@ -735,21 +780,22 @@ def _process_observations( elif resetted_obs != ASYNC_RESET_RETURN: # Creates a new episode if this is not async return. # If reset is async, we will get its result in some future poll - episode = active_episodes[env_id] + episode: MultiAgentEpisode = active_episodes[env_id] if observation_fn: - resetted_obs = observation_fn( + resetted_obs: Dict[AgentID, EnvObsType] = observation_fn( agent_obs=resetted_obs, worker=worker, base_env=base_env, policies=policies, episode=episode) + # type: AgentID, EnvObsType for agent_id, raw_obs in resetted_obs.items(): - policy_id = episode.policy_for(agent_id) - policy = _get_or_raise(policies, policy_id) - prep_obs = _get_or_raise(preprocessors, - policy_id).transform(raw_obs) - filtered_obs = _get_or_raise(obs_filters, - policy_id)(prep_obs) + policy_id: PolicyID = episode.policy_for(agent_id) + policy: Policy = _get_or_raise(policies, policy_id) + prep_obs: EnvObsType = _get_or_raise( + preprocessors, policy_id).transform(raw_obs) + filtered_obs: EnvObsType = _get_or_raise( + obs_filters, policy_id)(prep_obs) episode._set_last_observation(agent_id, filtered_obs) to_eval[policy_id].append( PolicyEvalData( @@ -763,27 +809,33 @@ def _process_observations( return active_envs, to_eval, outputs -def _do_policy_eval(*, to_eval, policies, active_episodes, tf_sess=None): +def _do_policy_eval( + *, + to_eval: Dict[PolicyID, List[PolicyEvalData]], + policies: Dict[PolicyID, Policy], + active_episodes: Dict[str, MultiAgentEpisode], + tf_sess=None +) -> Dict[PolicyID, Tuple[TensorStructType, StateBatch, dict]]: """Call compute_actions on collected episode/model data to get next action. Args: tf_sess (Optional[tf.Session]): Optional tensorflow session to use for batching TF policy evaluations. - to_eval (Dict[str,List[PolicyEvalData]]): Mapping of policy IDs to - lists of PolicyEvalData objects. - policies (Dict[str,Policy]): Mapping from policy ID to Policy obj. - active_episodes (defaultdict[str,MultiAgentEpisode]): Mapping from + to_eval (Dict[PolicyID, List[PolicyEvalData]]): Mapping of policy IDs + to lists of PolicyEvalData objects. + policies (Dict[PolicyID, Policy]): Mapping from policy ID to Policy. + active_episodes (Dict[str, MultiAgentEpisode]): Mapping from episode ID to currently ongoing MultiAgentEpisode object. Returns: eval_results: dict of policy to compute_action() outputs. """ - eval_results = {} + eval_results: Dict[PolicyID, TensorStructType] = {} if tf_sess: builder = TFRunBuilder(tf_sess, "policy_eval") - pending_fetches = {} + pending_fetches: Dict[PolicyID, Any] = {} else: builder = None @@ -791,16 +843,17 @@ def _do_policy_eval(*, to_eval, policies, active_episodes, tf_sess=None): logger.info("Inputs to compute_actions():\n\n{}\n".format( summarize(to_eval))) + # type: PolicyID, PolicyEvalData for policy_id, eval_data in to_eval.items(): - rnn_in = [t.rnn_state for t in eval_data] - policy = _get_or_raise(policies, policy_id) + rnn_in: List[List[Any]] = [t.rnn_state for t in eval_data] + policy: Policy = _get_or_raise(policies, policy_id) # If tf (non eager) AND TFPolicy's compute_action method has not been # overridden -> Use `policy._build_compute_actions()`. if builder and (policy.compute_actions.__code__ is TFPolicy.compute_actions.__code__): - obs_batch = [t.obs for t in eval_data] - state_batches = _to_column_format(rnn_in) + obs_batch: List[EnvObsType] = [t.obs for t in eval_data] + state_batches: StateBatch = _to_column_format(rnn_in) # TODO(ekl): how can we make info batch available to TF code? prev_action_batch = [t.prev_action for t in eval_data] prev_reward_batch = [t.prev_reward for t in eval_data] @@ -813,7 +866,7 @@ def _do_policy_eval(*, to_eval, policies, active_episodes, tf_sess=None): prev_reward_batch=prev_reward_batch, timestep=policy.global_timestep) else: - rnn_in_cols = [ + rnn_in_cols: StateBatch = [ np.stack([row[i] for row in rnn_in]) for i in range(len(rnn_in[0])) ] @@ -826,6 +879,7 @@ def _do_policy_eval(*, to_eval, policies, active_episodes, tf_sess=None): episodes=[active_episodes[t.env_id] for t in eval_data], timestep=policy.global_timestep) if builder: + # type: PolicyID, Tuple[TensorStructType, StateBatch, dict] for pid, v in pending_fetches.items(): eval_results[pid] = builder.get(v) @@ -836,25 +890,28 @@ def _do_policy_eval(*, to_eval, policies, active_episodes, tf_sess=None): return eval_results -def _process_policy_eval_results(*, to_eval, eval_results, active_episodes, - active_envs, off_policy_actions, policies, - clip_actions): +def _process_policy_eval_results( + *, to_eval: Dict[PolicyID, List[PolicyEvalData]], eval_results: Dict[ + PolicyID, Tuple[TensorStructType, StateBatch, dict]], + active_episodes: Dict[str, MultiAgentEpisode], active_envs: Set[int], + off_policy_actions: MultiEnvDict, policies: Dict[PolicyID, Policy], + clip_actions: bool) -> Dict[EnvID, Dict[AgentID, EnvActionType]]: """Process the output of policy neural network evaluation. Records policy evaluation results into the given episode objects and returns replies to send back to agents in the env. Args: - to_eval (Dict[str,List[PolicyEvalData]]): Mapping of policy IDs to - lists of PolicyEvalData objects. - eval_results (Dict[str,List]): Mapping of policy IDs to list of + to_eval (Dict[PolicyID, List[PolicyEvalData]]): Mapping of policy IDs + to lists of PolicyEvalData objects. + eval_results (Dict[PolicyID, List]): Mapping of policy IDs to list of actions, rnn-out states, extra-action-fetches dicts. - active_episodes (defaultdict[str,MultiAgentEpisode]): Mapping from + active_episodes (Dict[str, MultiAgentEpisode]): Mapping from episode ID to currently ongoing MultiAgentEpisode object. active_envs (Set[int]): Set of non-terminated env ids. off_policy_actions (dict): Doubly keyed dict of env-ids -> agent ids -> off-policy-action, returned by a `BaseEnv.poll()` call. - policies (Dict[str,Policy]): Mapping from policy ID to Policy obj. + policies (Dict[PolicyID, Policy]): Mapping from policy ID to Policy. clip_actions (bool): Whether to clip actions to the action space's bounds. @@ -862,16 +919,21 @@ def _process_policy_eval_results(*, to_eval, eval_results, active_episodes, actions_to_send: Nested dict of env id -> agent id -> agent replies. """ - actions_to_send = defaultdict(dict) + actions_to_send: Dict[EnvID, Dict[AgentID, EnvActionType]] = \ + defaultdict(dict) + + # type: int for env_id in active_envs: actions_to_send[env_id] = {} # at minimum send empty dict + # type: PolicyID, List[PolicyEvalData] for policy_id, eval_data in to_eval.items(): - rnn_in_cols = _to_column_format([t.rnn_state for t in eval_data]) + rnn_in_cols: StateBatch = _to_column_format( + [t.rnn_state for t in eval_data]) - actions = eval_results[policy_id][0] - rnn_out_cols = eval_results[policy_id][1] - pi_info_cols = eval_results[policy_id][2] + actions: TensorStructType = eval_results[policy_id][0] + rnn_out_cols: StateBatch = eval_results[policy_id][1] + pi_info_cols: dict = eval_results[policy_id][2] # In case actions is a list (representing the 0th dim of a batch of # primitive actions), try to convert it first. @@ -887,12 +949,13 @@ def _process_policy_eval_results(*, to_eval, eval_results, active_episodes, for f_i, column in enumerate(rnn_out_cols): pi_info_cols["state_out_{}".format(f_i)] = column - policy = _get_or_raise(policies, policy_id) + policy: Policy = _get_or_raise(policies, policy_id) # Split action-component batches into single action rows. - actions = unbatch(actions) + actions: List[EnvActionType] = unbatch(actions) + # type: int, EnvActionType for i, action in enumerate(actions): - env_id = eval_data[i].env_id - agent_id = eval_data[i].agent_id + env_id: int = eval_data[i].env_id + agent_id: AgentID = eval_data[i].agent_id # Clip if necessary. if clip_actions: clipped_action = clip_action(action, @@ -900,7 +963,7 @@ def _process_policy_eval_results(*, to_eval, eval_results, active_episodes, else: clipped_action = action actions_to_send[env_id][agent_id] = clipped_action - episode = active_episodes[env_id] + episode: MultiAgentEpisode = active_episodes[env_id] episode._set_rnn_state(agent_id, [c[i] for c in rnn_out_cols]) episode._set_last_pi_info( agent_id, {k: v[i] @@ -915,7 +978,7 @@ def _process_policy_eval_results(*, to_eval, eval_results, active_episodes, return actions_to_send -def _fetch_atari_metrics(base_env): +def _fetch_atari_metrics(base_env: BaseEnv) -> List[RolloutMetrics]: """Atari games have multiple logical episodes, one per life. However, for metrics reporting we count full episodes, all lives included. @@ -933,12 +996,13 @@ def _fetch_atari_metrics(base_env): return atari_out -def _to_column_format(rnn_state_rows): +def _to_column_format(rnn_state_rows: List[List[Any]]) -> StateBatch: num_cols = len(rnn_state_rows[0]) return [[row[i] for row in rnn_state_rows] for i in range(num_cols)] -def _get_or_raise(mapping, policy_id): +def _get_or_raise(mapping: Dict[PolicyID, Policy], + policy_id: PolicyID) -> Policy: """Returns a Policy object under key `policy_id` in `mapping`. Args: diff --git a/rllib/evaluation/worker_set.py b/rllib/evaluation/worker_set.py index 5e52a0fc1..ed79b6444 100644 --- a/rllib/evaluation/worker_set.py +++ b/rllib/evaluation/worker_set.py @@ -1,5 +1,6 @@ import logging from types import FunctionType +from typing import TypeVar, Callable, List, Union import ray from ray.rllib.utils.annotations import DeveloperAPI @@ -7,13 +8,19 @@ from ray.rllib.evaluation.rollout_worker import RolloutWorker, \ _validate_multiagent_config from ray.rllib.offline import NoopOutput, JsonReader, MixedInput, JsonWriter, \ ShuffledInput +from ray.rllib.env.env_context import EnvContext +from ray.rllib.policy import Policy from ray.rllib.utils import merge_dicts from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.types import PolicyID, TrainerConfigDict, EnvType tf = try_import_tf() logger = logging.getLogger(__name__) +# Generic type var for foreach_* methods. +T = TypeVar("T") + @DeveloperAPI class WorkerSet: @@ -23,12 +30,12 @@ class WorkerSet: """ def __init__(self, - env_creator, - policy, - trainer_config=None, - num_workers=0, - logdir=None, - _setup=True): + env_creator: Callable[[EnvContext], EnvType], + policy: type, + trainer_config: TrainerConfigDict = None, + num_workers: int = 0, + logdir: str = None, + _setup: bool = True): """Create a new WorkerSet and initialize its workers. Arguments: @@ -63,22 +70,22 @@ class WorkerSet: self._remote_workers = [] self.add_workers(num_workers) - def local_worker(self): + def local_worker(self) -> RolloutWorker: """Return the local rollout worker.""" return self._local_worker - def remote_workers(self): + def remote_workers(self) -> List["ActorHandle"]: """Return a list of remote rollout workers.""" return self._remote_workers - def sync_weights(self): + def sync_weights(self) -> None: """Syncs weights of remote workers with the local worker.""" if self.remote_workers(): weights = ray.put(self.local_worker().get_weights()) for e in self.remote_workers(): e.set_weights.remote(weights) - def add_workers(self, num_workers): + def add_workers(self, num_workers: int) -> None: """Creates and add a number of remote workers to this worker set. Args: @@ -99,11 +106,11 @@ class WorkerSet: self._remote_config) for i in range(num_workers) ]) - def reset(self, new_remote_workers): + def reset(self, new_remote_workers: List["ActorHandle"]) -> None: """Called to change the set of remote workers.""" self._remote_workers = new_remote_workers - def stop(self): + def stop(self) -> None: """Stop all rollout workers.""" self.local_worker().stop() for w in self.remote_workers(): @@ -111,7 +118,7 @@ class WorkerSet: w.__ray_terminate__.remote() @DeveloperAPI - def foreach_worker(self, func): + def foreach_worker(self, func: Callable[[RolloutWorker], T]) -> List[T]: """Apply the given function to each worker instance.""" local_result = [func(self.local_worker())] @@ -120,7 +127,8 @@ class WorkerSet: return local_result + remote_results @DeveloperAPI - def foreach_worker_with_index(self, func): + def foreach_worker_with_index( + self, func: Callable[[RolloutWorker, int], T]) -> List[T]: """Apply the given function to each worker instance. The index will be passed as the second arg to the given function. @@ -133,7 +141,7 @@ class WorkerSet: return local_result + remote_results @DeveloperAPI - def foreach_policy(self, func): + def foreach_policy(self, func: Callable[[Policy, PolicyID], T]) -> List[T]: """Apply the given function to each worker's (policy, policy_id) tuple. Args: @@ -153,12 +161,13 @@ class WorkerSet: return local_results + remote_results @DeveloperAPI - def trainable_policies(self): + def trainable_policies(self) -> List[PolicyID]: """Return the list of trainable policy ids.""" return self.local_worker().foreach_trainable_policy(lambda _, pid: pid) @DeveloperAPI - def foreach_trainable_policy(self, func): + def foreach_trainable_policy( + self, func: Callable[[Policy, PolicyID], T]) -> List[T]: """Apply `func` to all workers' Policies iff in `policies_to_train`. Args: @@ -179,13 +188,17 @@ class WorkerSet: return local_results + remote_results @staticmethod - def _from_existing(local_worker, remote_workers=None): + def _from_existing(local_worker: RolloutWorker, + remote_workers: List["ActorHandle"] = None): workers = WorkerSet(None, None, {}, _setup=False) workers._local_worker = local_worker workers._remote_workers = remote_workers or [] return workers - def _make_worker(self, cls, env_creator, policy, worker_index, config): + def _make_worker( + self, cls: Callable, env_creator: Callable[[EnvContext], EnvType], + policy: Policy, worker_index: int, + config: TrainerConfigDict) -> Union[RolloutWorker, "ActorHandle"]: def session_creator(): logger.debug("Creating TF session {}".format( config["tf_session_args"])) diff --git a/rllib/execution/common.py b/rllib/execution/common.py index a7741e832..f63ac7475 100644 --- a/rllib/execution/common.py +++ b/rllib/execution/common.py @@ -1,5 +1,3 @@ -from typing import Union - from ray.util.iter import LocalIterator from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch @@ -23,14 +21,11 @@ LOAD_BATCH_TIMER = "load" # Instant metrics (keys for metrics.info). LEARNER_INFO = "learner" -# Type aliases. -GradientType = dict -SampleBatchType = Union[SampleBatch, MultiAgentBatch] - # Asserts that an object is a type of SampleBatch. def _check_sample_batch_type(batch): - if not isinstance(batch, SampleBatchType.__args__): + if not isinstance(batch, SampleBatch) and not isinstance( + batch, MultiAgentBatch): raise ValueError("Expected either SampleBatch or MultiAgentBatch, " "got {}: {}".format(type(batch), batch)) diff --git a/rllib/execution/replay_buffer.py b/rllib/execution/replay_buffer.py index 655190e7c..dba78e790 100644 --- a/rllib/execution/replay_buffer.py +++ b/rllib/execution/replay_buffer.py @@ -6,7 +6,6 @@ import random from typing import List import ray -from ray.rllib.execution.common import SampleBatchType from ray.rllib.execution.segment_tree import SumSegmentTree, MinSegmentTree from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch, \ DEFAULT_POLICY_ID @@ -14,6 +13,7 @@ from ray.rllib.utils.annotations import DeveloperAPI from ray.util.iter import ParallelIteratorWorker from ray.rllib.utils.timer import TimerStat from ray.rllib.utils.window_stat import WindowStat +from ray.rllib.utils.types import SampleBatchType # Constant that represents all policies in lockstep replay mode. _ALL_POLICIES = "__all__" diff --git a/rllib/execution/replay_ops.py b/rllib/execution/replay_ops.py index 6a5bf96dd..c5245e78f 100644 --- a/rllib/execution/replay_ops.py +++ b/rllib/execution/replay_ops.py @@ -4,8 +4,9 @@ import random from ray.util.iter import from_actors, LocalIterator, _NextValueNotReady from ray.util.iter_metrics import SharedMetrics from ray.rllib.execution.replay_buffer import LocalReplayBuffer -from ray.rllib.execution.common import SampleBatchType, \ +from ray.rllib.execution.common import \ STEPS_SAMPLED_COUNTER, _get_shared_metrics +from ray.rllib.utils.types import SampleBatchType class StoreToReplayBuffer: diff --git a/rllib/execution/rollout_ops.py b/rllib/execution/rollout_ops.py index a60b9b734..c199e15b9 100644 --- a/rllib/execution/rollout_ops.py +++ b/rllib/execution/rollout_ops.py @@ -7,13 +7,13 @@ from ray.util.iter_metrics import SharedMetrics from ray.rllib.evaluation.metrics import get_learner_stats from ray.rllib.evaluation.rollout_worker import get_global_worker from ray.rllib.evaluation.worker_set import WorkerSet -from ray.rllib.execution.common import GradientType, SampleBatchType, \ - STEPS_SAMPLED_COUNTER, LEARNER_INFO, SAMPLE_TIMER, \ - GRAD_WAIT_TIMER, _check_sample_batch_type, _get_shared_metrics -from ray.rllib.policy.policy import PolicyID +from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER, LEARNER_INFO, \ + SAMPLE_TIMER, GRAD_WAIT_TIMER, _check_sample_batch_type, \ + _get_shared_metrics from ray.rllib.policy.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \ MultiAgentBatch from ray.rllib.utils.sgd import standardized +from ray.rllib.utils.types import PolicyID, SampleBatchType, ModelGradients logger = logging.getLogger(__name__) @@ -91,7 +91,7 @@ def ParallelRollouts(workers: WorkerSet, *, mode="bulk_sync", def AsyncGradients( - workers: WorkerSet) -> LocalIterator[Tuple[GradientType, int]]: + workers: WorkerSet) -> LocalIterator[Tuple[ModelGradients, int]]: """Operator to compute gradients in parallel from rollout workers. Arguments: diff --git a/rllib/execution/train_ops.py b/rllib/execution/train_ops.py index 860337cb0..f4e794740 100644 --- a/rllib/execution/train_ops.py +++ b/rllib/execution/train_ops.py @@ -7,18 +7,18 @@ from typing import List import ray from ray.rllib.evaluation.metrics import get_learner_stats, LEARNER_STATS_KEY from ray.rllib.evaluation.worker_set import WorkerSet -from ray.rllib.execution.common import SampleBatchType, \ +from ray.rllib.execution.common import \ STEPS_SAMPLED_COUNTER, STEPS_TRAINED_COUNTER, LEARNER_INFO, \ APPLY_GRADS_TIMER, COMPUTE_GRADS_TIMER, WORKER_UPDATE_TIMER, \ LEARN_ON_BATCH_TIMER, LOAD_BATCH_TIMER, LAST_TARGET_UPDATE_TS, \ NUM_TARGET_UPDATES, _get_global_vars, _check_sample_batch_type, \ _get_shared_metrics from ray.rllib.execution.multi_gpu_impl import LocalSyncParallelOptimizer -from ray.rllib.policy.policy import PolicyID from ray.rllib.policy.sample_batch import SampleBatch, DEFAULT_POLICY_ID, \ MultiAgentBatch from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.sgd import do_minibatch_sgd, averaged +from ray.rllib.utils.types import PolicyID, SampleBatchType tf = try_import_tf() diff --git a/rllib/execution/tree_agg.py b/rllib/execution/tree_agg.py index bb8442242..b7d07850b 100644 --- a/rllib/execution/tree_agg.py +++ b/rllib/execution/tree_agg.py @@ -4,12 +4,13 @@ from typing import List import ray from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER, \ - SampleBatchType, _get_shared_metrics + _get_shared_metrics from ray.rllib.execution.replay_ops import MixInReplay from ray.rllib.execution.rollout_ops import ParallelRollouts, ConcatBatches from ray.rllib.utils.actors import create_colocated from ray.util.iter import ParallelIterator, ParallelIteratorWorker, \ from_actors +from ray.rllib.utils.types import SampleBatchType logger = logging.getLogger(__name__) diff --git a/rllib/policy/__init__.py b/rllib/policy/__init__.py index c419ee35e..348fe187d 100644 --- a/rllib/policy/__init__.py +++ b/rllib/policy/__init__.py @@ -1,12 +1,10 @@ -from ray.rllib.policy.policy import Policy, PolicyID, AgentID +from ray.rllib.policy.policy import Policy from ray.rllib.policy.torch_policy import TorchPolicy from ray.rllib.policy.tf_policy import TFPolicy from ray.rllib.policy.torch_policy_template import build_torch_policy from ray.rllib.policy.tf_policy_template import build_tf_policy __all__ = [ - "AgentID", - "PolicyID", "Policy", "TFPolicy", "TorchPolicy", diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 63a4a7401..d0e5dd0d4 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -1,7 +1,6 @@ from abc import ABCMeta, abstractmethod import gym import numpy as np -from typing import Any from ray.rllib.utils import try_import_tree from ray.rllib.utils.annotations import DeveloperAPI @@ -18,12 +17,6 @@ tree = try_import_tree() # `grad_info` dict returned by learn_on_batch() / compute_grads() via this key. LEARNER_STATS_KEY = "learner_stats" -# Represents a generic identifier for an agent (e.g., "agent1"). -AgentID = Any - -# Represents a generic identifier for a policy (e.g., "pol1"). -PolicyID = str - @DeveloperAPI class Policy(metaclass=ABCMeta): diff --git a/rllib/utils/framework.py b/rllib/utils/framework.py index c0434126c..014d4d7d6 100644 --- a/rllib/utils/framework.py +++ b/rllib/utils/framework.py @@ -6,9 +6,11 @@ from typing import Any, Union logger = logging.getLogger(__name__) # Represents a generic tensor type. +# TODO(ekl) this is duplicated in types.py TensorType = Any # Either a plain tensor, or a dict or tuple of tensors (or StructTensors). +# TODO(ekl) this is duplicated in types.py TensorStructType = Union[TensorType, dict, tuple] diff --git a/rllib/utils/types.py b/rllib/utils/types.py new file mode 100644 index 000000000..a43b10497 --- /dev/null +++ b/rllib/utils/types.py @@ -0,0 +1,79 @@ +from typing import Any, Dict, Union, Tuple +import gym + +# Represents a fully filled out config of a Trainer class. +TrainerConfigDict = dict + +# A trainer config dict that only has overrides. It needs to be combined with +# the default trainer config to be used. +PartialTrainerConfigDict = dict + +# Represents the env_config sub-dict of the trainer config that is passed to +# the env constructor. +EnvConfigDict = dict + +# Represents the model config sub-dict of the trainer config that is passed to +# the model catalog. +ModelConfigDict = dict + +# Represents a BaseEnv, MultiAgentEnv, ExternalEnv, ExternalMultiAgentEnv, +# VectorEnv, or gym.Env. +EnvType = Any + +# Represents a generic identifier for an agent (e.g., "agent1"). +AgentID = Any + +# Represents a generic identifier for a policy (e.g., "pol1"). +PolicyID = str + +# Type of the config["multiagent"]["policies"] dict for multi-agent training. +MultiAgentPolicyConfigDict = Dict[PolicyID, Tuple[type, gym.Space, gym.Space, + PartialTrainerConfigDict]] + +# Represents an environment id. +EnvID = int + +# A dict keyed by agent ids, e.g. {"agent-1": value}. +MultiAgentDict = Dict[AgentID, Any] + +# A dict keyed by env ids that contain further nested dictionaries keyed by +# agent ids. e.g., {"env-1": {"agent-1": value}}. +MultiEnvDict = Dict[EnvID, MultiAgentDict] + +# Represents an observation returned from the env. +EnvObsType = Any + +# Represents an action passed to the env. +EnvActionType = Any + +# Info dictionary returned by calling step() on gym envs. Commonly empty dict. +EnvInfoDict = dict + +# Represents the result dict returned by Trainer.train(). +ResultDict = dict + +# Dict of tensors returned by compute gradients on the policy, e.g., +# {"td_error": [...], "learner_stats": {"vf_loss": ..., ...}}, for multi-agent, +# {"policy1": {"learner_stats": ..., }, "policy2": ...}. +GradInfoDict = dict + +# Dict of learner stats returned by compute gradients on the policy, e.g., +# {"vf_loss": ..., ...}. This will always be nested under the "learner_stats" +# key(s) of a GradInfoDict. In the multi-agent case, this will be keyed by +# policy id. +LearnerStatsDict = dict + +# Type of dict returned by compute_gradients() representing model gradients. +ModelGradients = dict + +# Type of dict returned by get_weights() representing model weights. +ModelWeights = dict + +# Some kind of sample batch. +SampleBatchType = Union["SampleBatch", "MultiAgentBatch"] + +# 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]