wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-13 08:10:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rllib] Add type annotations for evaluation/, env/ packages (#9003)

Browse Source
This commit is contained in:
Eric Liang
2020-06-19 13:09:05 -07:00
committed by GitHub
parent f43cad6371
commit 1e0e1a45e6
34 changed files with 840 additions and 500 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/agents/callbacks.py
+2 -1
View File
@@ -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
rllib/agents/dqn/apex.py
+2 -1
View File
@@ -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__
rllib/agents/registry.py
+2 -2
View File
@@ -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:
rllib/agents/trainer.py
+32 -21
View File
@@ -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.
rllib/agents/trainer_template.py
+16 -11
View File
@@ -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:
rllib/env/__init__.py
Vendored
+2 -5
View File
@@ -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",
]
rllib/env/base_env.py
Vendored
+51 -29
View File
@@ -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
rllib/env/env_context.py
Vendored
+10 -5
View File
@@ -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,
rllib/env/external_env.py
Vendored
+24 -11
View File
@@ -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
rllib/env/external_multi_agent_env.py
Vendored
+20 -9
View File
@@ -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:
rllib/env/multi_agent_env.py
Vendored
+13 -3
View File
@@ -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
rllib/env/policy_client.py
Vendored
+29 -15
View File
@@ -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)
rllib/env/policy_server_input.py
Vendored
+2 -2
View File
@@ -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())
rllib/env/remote_vector_env.py
Vendored
+19 -6
View File
@@ -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()
rllib/env/unity3d_env.py
Vendored
+20 -14
View File
@@ -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, )),
]),
rllib/env/vector_env.py
Vendored
+18 -11
View File
@@ -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:
rllib/evaluation/episode.py
+57 -36
View File
@@ -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,
rllib/evaluation/metrics.py
+19 -11
View File
@@ -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 = [], []
rllib/evaluation/observation_function.py
+2 -1
View File
@@ -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:
rllib/evaluation/postprocessing.py
+7 -7
View File
@@ -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.
rllib/evaluation/rollout_worker.py
+151 -114
View File
@@ -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
rllib/evaluation/sample_batch_builder.py
+23 -12
View File
@@ -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
rllib/evaluation/sampler.py
+191 -127
View File
@@ -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:
rllib/evaluation/worker_set.py
+32 -19
View File
@@ -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"]))
rllib/execution/common.py
+2 -7
View File
@@ -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))
rllib/execution/replay_buffer.py
+1 -1
View File
@@ -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__"
rllib/execution/replay_ops.py
+2 -1
View File
@@ -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:
rllib/execution/rollout_ops.py
+5 -5
View File
@@ -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:
rllib/execution/train_ops.py
+2 -2
View File
@@ -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()
rllib/execution/tree_agg.py
+2 -1
View File
@@ -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__)
rllib/policy/__init__.py
+1 -3
View File
@@ -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",
rllib/policy/policy.py
-7
View File
@@ -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):
rllib/utils/framework.py
+2
View File
@@ -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]
rllib/utils/types.py
+79
View File
@@ -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]